from viam.components.base import Base
from viam.components.camera import Camera
from viam.errors import ResourceNotFoundError

class Rover:
    def __init__(self, robot):
        self.robot = robot
        try:
            self.camera = Camera.from_robot(robot=robot, name="cam")
        except ResourceNotFoundError:
            self.camera = Camera.from_robot(robot=robot, name="cam2")
        self.base = Base.from_robot(robot, "viam_base")

    async def capture_image(self, image_count):
        # take a picture and send it to ML model for better accurcy
        for i in range(image_count):
            frame = await self.camera.get_image()
            if frame is not None:
                # Convert frame to a format that can be saved
                file_path = f"../../images/test_image_{i}.jpg"
                await self.base.spin(velocity=100, angle=63)
                with open(file_path, "wb") as f:
                    f.write(frame.data)
            else:
                print("Error cannot capture picture")
        return

    def __del__(self):
        self.robot.close()

backend/src/Rover.py
prachee-n16/pic-perfect
# Import the necessary libraries
import cv2
import csv
import os
import matplotlib.pyplot as plt
from skimage.exposure import is_low_contrast
import numpy as np

# create an array of all the input images
fileNames = []
modifyOrigImages = []
openedImages = []
imageFlags = []
imagesPath = "../../client/src/captures/"
threshold = 100
fileNames = os.listdir(imagesPath)


def variance_of_laplacian(img2):
    # compute the Laplacian of the image and then return the focus
    # measure, which is simply the variance of the Laplacian
    gray = cv2.cvtColor(img2, cv2.COLOR_RGB2BGR)
    return cv2.Laplacian(gray, cv2.CV_64F).var()


def BGR2RGB(BGR_img):
    # turning BGR pixel color to RGB
    rgb_image = cv2.cvtColor(BGR_img, cv2.COLOR_BGR2RGB)
    return rgb_image


# read the input images into input array
def readImages(index):
    for file in fileNames:
        openedImages.append(cv2.imread(file))


def evaluateImages(length):
    fileNames = os.listdir(imagesPath)
    print(fileNames)
    with open("../../client/imagesInfo.txt", "w") as csvfile:
        csvwriter = csv.writer(csvfile)
        # criteria: blur, contrast, sharpness, and color balance
        for i in range(length):
            fileName = fileNames[i]
            currentImgFlags = []
            img = cv2.imread(imagesPath + fileName)
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            # populate image flags
            # array format: [ image file name, [flags], corrected file name], where the indices of flags are:
            # index 0 : contrast, 1 for too high, 0 for too low
            # index 1 : blurriness, 1 for too blurry, 0 for clear

            # file gets corrected if the contrast is too low/high
            if is_low_contrast(gray, 0.35):
                currentImgFlags.append(1)
            else:
                currentImgFlags.append(0)
            # file gets corrected if the blur is too much / adds sharpenning
            # if the file has a variance less than threshold, it is blurry. fix sharpness or blurriness accordingly
            if variance_of_laplacian(img) < threshold:
                currentImgFlags.append(1)
            else:
                currentImgFlags.append(0)

            # Implement fixes:
            if currentImgFlags[0] == 1:
                # image is high contrast, balance it out
                contrast = 2.3
                image2 = cv2.addWeighted(
                    img, contrast, np.zeros(img.shape, img.dtype), 0, 40
                )
            else:
                contrast = 1.0
                image2 = cv2.addWeighted(
                    img, contrast, np.zeros(img.shape, img.dtype), 0, 20
                )

            if currentImgFlags[1] == 1:
                # image is too blurry, apply sharpness
                kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
                image2 = cv2.filter2D(img, ddepth=-1, kernel=kernel)
            else:
                # image is too sharp, apply blur
                image2 = cv2.GaussianBlur(img, (3, 3), 0)

            # always apply color correction to balance colors:
            alpha = 1.5  # Contrast control (1.0 means no change)
            beta = 30  # Brightness control (0 means no change)
            image2 = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)
            hsv_image = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
            saturation_factor = 1.5  # Adjust this value to change saturation
            hsv_image[:, :, 1] = np.clip(
                hsv_image[:, :, 1] * saturation_factor, 0, 255)
            image2 = cv2.cvtColor(hsv_image, cv2.COLOR_HSV2BGR)
            # Define a matrix for color correction (can be a custom matrix)
            color_correction_matrix = np.array(
                [[1.2, 0.1, 0.1], [0.0, 1.3, 0.2], [0.1, 0.2, 1.4]]
            )

            image2 = cv2.transform(img, color_correction_matrix)
            correctedFile = "../../client/src/corrected/" + \
                fileName[:-5] + "_enhanced.jpeg" + "\n"
            cv2.imwrite(correctedFile, image2)
            imageFlags.append(
                ["../../client/src/captures/" + fileName, correctedFile + "\n"])
            cv2.waitKey(0)
            currentImgFlags = []
            # closing all open windows
            cv2.destroyAllWindows()
        csvwriter.writerows(imageFlags)

backend/src/filters.py
prachee-n16/pic-perfect
import asyncio
import cv2
import json
import mediapipe as mp
import numpy as np
import random
import time

from filters import evaluateImages
from datetime import datetime
from tensorflow.keras.models import load_model
from viam.robot.client import RobotClient
from viam.rpc.dial import Credentials, DialOptions

from Rover import Rover

VELOCITY = 1000


async def connect():
    with open("../key.json", "r") as f:
        payload = json.load(f)["payload"]
    creds = Credentials(type="robot-location-secret", payload=payload)
    opts = RobotClient.Options(
        refresh_interval=0, dial_options=DialOptions(credentials=creds)
    )
    return await RobotClient.at_address("htn-main.3o4e4wpz0f.viam.cloud", opts)


def leftOrRight(face_cascade, midpoint, frame_size, withinRange):
    min_face_size = 50  # Adjust this based on your requirements

    largest_area = 0
    largest = ()
    image = cv2.imread("img.jpeg")
    # Convert into grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # Detect faces
    faces = face_cascade.detectMultiScale(gray, 1.1, 4)
    if len(faces) == 0:
        print("nothing detected")
        return [-1]

    # Draw rectangle around the faces
    for face in faces:
        (x, y, w, h) = face
        a = w * h
        if a > largest_area:
            largest_area = a
            largest = face
        cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 0), 2)

    (x, y, w, h) = largest
    # Need to test this stop logic here, will this even work?
    if w >= min_face_size and h >= min_face_size:
        x_within_range = x >= 0.15 * \
            frame_size[0] and x + w <= 0.85 * frame_size[0]
        y_within_range = y >= 0.15 * \
            frame_size[1] and y + h <= 0.6 * frame_size[1]
        # print(x_within_range, y_within_range)
        # print(0.01 *
        #       frame_size[0])
        # print(x + w >= 0.99 * frame_size[0])
        # print(y <= 0.01 *
        #       frame_size[1])
        # print(y + h >= 0.99 * frame_size[1])
        if x_within_range and y_within_range:
            withinRange[0] = 1

    cx = (x + (x + w)) / 2
    if cx < (midpoint - (midpoint / 6)):
        return [0]
    if cx > (midpoint + (midpoint / 6)):
        return [2]
    elif y <= 0.15 * frame_size[1]:
        return [3]
    elif y + h >= 0.60 * frame_size[1]:
        return [1]
    return [-1]


async def main():
    robot = await connect()

    # Create instance of the rover class
    my_rover = Rover(robot)

    print("Robot Connected!")
    print("Resources:")
    print(robot.resource_names)

    face_cascade = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")

    # Use MediaPipe to draw the hand frameword over the hands it identifies in synchronous
    hands_module = mp.solutions.hands
    draw_module = mp.solutions.drawing_utils

    # Use OpenCV to create a Video stream and add some values
    camera = my_rover.camera

    # intiate model
    model = load_model("../src/mp_hand_gesture")
    file = open("gesture.names", "r")
    classNames = file.read().split("\n")
    file.close()

    # We add confidence values and extra settings to MediaPipe hand tracking. Since this is a live video stream and not static
    # image mode, confidence values in regards to overall detection and tracking and we will only let one hands be tracked at a time
    with hands_module.Hands(
        static_image_mode=False,
        min_detection_confidence=0.7,
        min_tracking_confidence=0.7,
        max_num_hands=1,
    ) as hands:
        while True:
            pastMove = ""

            frame = await camera.get_image()
            if frame is not None:
                file_path = "img.jpeg"
                with open(file_path, "wb") as f:
                    f.write(frame.data)
            else:
                print("Error capturing image.")
                break

            frame1 = await my_rover.camera.get_image(mime_type="image/jpeg")
            if frame1 is None:
                print("Error capturing image.")
                break

            # get a cv2 version of the frame
            frame = cv2.imread(file_path)

            # for balance between speed and accurate and identification
            refined_frame = cv2.resize(frame, (640, 480))
            x, y, c = refined_frame.shape

            # Produces the hand framework overlay ontop of the hand
            result = hands.process(cv2.cvtColor(
                refined_frame, cv2.COLOR_BGR2RGB))

            className = ""

            if result.multi_hand_landmarks:
                Landmarks = []
                for hand_landmarks in result.multi_hand_landmarks:
                    for lm in hand_landmarks.landmark:
                        lmx = int(lm.x * x)
                        lmy = int(lm.y * y)

                        Landmarks.append([lmx, lmy])

                    # Drawing landmarks on frames
                    draw_module.draw_landmarks(
                        refined_frame,
                        hand_landmarks,
                        hands_module.HAND_CONNECTIONS,
                    )

                    # Predict gesture
                    prediction = model.predict([Landmarks])
                    classID = np.argmax(prediction)
                    className = classNames[classID]

            withinRange = [0]
            takePicture = False

            if className == "thumbs up":
                print("Time to start capturing pictures!")
                takePicture = True
                withinRange[0] = 1

            if withinRange[0] != 1:
                answer = leftOrRight(
                    face_cascade, frame1.size[0] / 2, frame1.size, withinRange
                )
            if withinRange[0] != 1:
                if answer[0] == 0:
                    # print("left")
                    # CCW is positive
                    pastMove = "left"
                    await my_rover.base.spin(angle=5, velocity=VELOCITY)
                    await my_rover.base.move_straight(distance=50, velocity=VELOCITY)
                elif answer[0] == 1:
                    # print("center")
                    pastMove = "center"
                    await my_rover.base.move_straight(distance=100, velocity=VELOCITY)
                elif answer[0] == 2:
                    # print("right")
                    pastMove = "right"
                    await my_rover.base.spin(angle=-5, velocity=VELOCITY)
                    await my_rover.base.move_straight(distance=50, velocity=VELOCITY)
                elif answer[0] == 3:
                    pastMove = "backward"
                    await my_rover.base.move_straight(distance=-100, velocity=VELOCITY)
                elif answer[0] == -1:
                    if pastMove == "":
                        print("Randomly moving to look for objects.")
                        # Generate random movements (example: forward, backward, left, right)
                        random_movement = random.choice(
                            ["forward", "backward", "left", "right"]
                        )
                        if random_movement == "forward":
                            await my_rover.base.move_straight(
                                distance=100, velocity=VELOCITY
                            )
                        elif random_movement == "backward":
                            await my_rover.base.move_straight(
                                distance=-100, velocity=VELOCITY
                            )
                        elif random_movement == "left":
                            await my_rover.base.spin(angle=5, velocity=VELOCITY)
                        elif random_movement == "right":
                            await my_rover.base.spin(angle=-5, velocity=VELOCITY)
                    elif pastMove == "right":
                        await my_rover.base.spin(angle=5, velocity=VELOCITY)
                        pastMove = ""
                    elif pastMove == "left":
                        await my_rover.base.spin(angle=-5, velocity=VELOCITY)
                        pastMove = ""
                    elif pastMove == "center":
                        await my_rover.base.move_straight(
                            distance=-100, velocity=VELOCITY
                        )
                        pastMove = ""
                    else:
                        await my_rover.base.move_straight(
                            distance=100, velocity=VELOCITY
                        )
                        pastMove = ""
            else:
                print("STOP!")
                await my_rover.base.stop()

            if takePicture:
                for _ in range(5):
                    time.sleep(1)
                    await my_rover.base.stop()
                for _ in range(3):
                    img = await my_rover.camera.get_image()
                    if img is not None:
                        current_date = datetime.now().strftime("%Y%m%d_%H%M%S")
                        with open(f"../../client/src/captures/capture_{current_date}.jpeg", "wb") as f:
                            f.write(img.data)
                    else:
                        print("Error capturing image.")
                        break
                    time.sleep(1)
                evaluateImages(3)

    # close the robot when you're done!
    del my_rover


if __name__ == "__main__":
    asyncio.run(main())

backend/src/main.py
prachee-n16/pic-perfect