Add apple

app.py

@@ -0,0 +1,36 @@
+import cv2
+import gradio as gr
+from predict_image import load_model, predict
+
+def predict_fn(image, model_name):
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    cv2.imwrite('./myimage.jpg', image)
+    # model for emotion classification
+    if model_name == 'EfficientNetB0':
+        model_name = 'effb0'
+    elif model_name == 'ResNet18':
+        model_name = 'res18'
+    else:
+        raise ValueError('Enter correct model_name')
+    
+    model = load_model(model_name)
+    
+    out = predict('./myimage.jpg', './result.jpg', model)
+    out = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
+    return out
+
+
+demo = gr.Interface(
+    fn=predict_fn,
+    inputs=[
+        gr.inputs.Image(label="Input Image"),
+        gr.Radio(['EfficientNetB0', 'ResNet18'], value='EfficientNetB0', label='Model Name')
+        ],
+    outputs=[
+        gr.inputs.Image(label="Prediction"),
+    ],
+    title="Emotion Recognition Demo",
+    description="Emotion Classification Model trained on FER Dataset"
+)
+
+demo.launch(debug=True)

face_module.py

@@ -0,0 +1,24 @@
+import mediapipe as mp
+mp_face_detection = mp.solutions.face_detection
+
+def get_face_coords(image):
+    with mp_face_detection.FaceDetection(
+        model_selection=1, min_detection_confidence=0.5) as face_detection:
+        #image = cv2.imread(file)
+        # Convert the BGR image to RGB and process it with MediaPipe Face Detection.
+        results = face_detection.process(image)
+        # Draw face detections of each face.
+        if not results.detections:
+            return False
+
+        # shape of image
+        h, w, _ = image.shape
+        
+        t = results.detections[0].location_data.relative_bounding_box
+        height = t.height * h
+        ymin = t.ymin * h
+        width = t.width * w
+        xmin = t.xmin * w
+        xmax = xmin + width
+        ymax = ymin + height
+        return int(xmin), int(ymin), int(xmax), int(ymax)

meter.py

@@ -0,0 +1,109 @@
+import cv2
+import numpy as np
+
+# storing settings for semicircle
+class SemiCircle:
+    def __init__(
+        self, thickness=10, color=(255, 0, 0), radius=100, 
+        center=(250, 250), angle=0, start_angle=180, end_angle=360
+        ):
+        self.thickness = thickness
+        self.color = color
+        self.radius = (radius, radius)
+        self.center = center
+        self.angle = angle
+        self.start_angle = start_angle
+        self.end_angle = end_angle
+
+class Line:
+     def __init__(self, thickness=2, color=(0, 0, 0)):
+        self.thickness = thickness
+        self.color = color
+
+def generate_points(radius, length, center, num_points):
+    # center points
+    cx, cy = center
+
+    # generating points on circle
+    outer_circle_points = [(radius * np.cos(i), radius * np.sin(i)) for i in np.linspace(np.pi, 2*np.pi, num_points)]
+
+    inner_radius = radius - length
+    inner_circle_points = [(inner_radius * np.cos(i), inner_radius * np.sin(i)) for i in np.linspace(np.pi, 2*np.pi, num_points)]
+
+
+    # genrating point for drawing line using cv2, start_points and end_points
+    start_points = [(int(cx + i), int(cy + j)) for i, j in outer_circle_points]
+    end_points = [(int(cx + i), int(cy + j)) for i, j in inner_circle_points]
+    return zip(start_points, end_points)
+
+class Meter:
+    def __init__(self, center, radius, circle_color):
+        self.center = center
+        self.radius = radius
+        self.circle_color = circle_color
+
+    def draw_meter(self, image, idx):
+        # drawing semicircle
+        circle = SemiCircle(center=self.center, radius=self.radius, color=self.circle_color)
+
+        cv2.ellipse(
+            image, circle.center, circle.radius, 
+            circle.angle, circle.start_angle, circle.end_angle, 
+            circle.color, circle.thickness
+        )
+
+        # drawing smaller fine lines
+        line = Line(thickness=circle.thickness//3)
+
+        for start, end in generate_points(self.radius - 10, self.radius * 0.05, self.center, 50):
+            cv2.line(image, start, end, line.color, line.thickness)
+
+        # drawing bigger lines
+        line2 = Line(thickness=circle.thickness//2, color=(238, 222, 23))
+
+        for start, end in generate_points(self.radius - 10, self.radius * 0.15, self.center, 10):
+            cv2.line(image, start, end, line2.color, line2.thickness)
+
+        # drawing needle anchor point
+        cv2.circle(image, circle.center, 15, (0, 0, 0), -1)
+
+        # emotion classes
+        emotions = ['neutral', 'happy', 'surprise', 'sad', 'angry', 
+                    "disgust", 'fear', 'contempt', 'unknown', 'NotFace']
+
+        # points where text will be written
+        pts = generate_points(self.radius*1.55, self.radius * 0.15, self.center, 10)
+        pts = list(pts)
+
+        for i, emot in enumerate(emotions):
+            x, y = pts[i][0]
+            x -= 20
+
+            color = (98, 65, 255) if (i == idx) else (144,238,144)
+            
+            cv2.putText(
+                image, emot, (x, y), 
+                cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2, 
+                cv2.LINE_AA
+            )           
+
+        # needle 12, 178, 33
+        line3 = Line(thickness=circle.thickness//2, color=(0, 0, 255))
+        pts2 = generate_points(self.radius*0.7, self.radius*0.7, self.center, 10)
+        pts2 = list(pts2)
+        start, end = pts2[idx]
+
+        cv2.line(image, start, end, line3.color, line3.thickness)
+
+if __name__ == '__main__':
+    image = np.ones((500, 500, 3))
+
+    meter = Meter((250, 250), 150, (80, 127, 255))
+
+    meter.draw_meter(image, 4)
+
+    cv2.imshow('image', image)
+
+    if cv2.waitKey(0) & 0xFF == 27:
+            pass
+    cv2.destroyAllWindows()

model.py

@@ -0,0 +1,25 @@
+import timm
+import torch.nn as nn
+from torchvision import models
+
+class ResnetModel(nn.Module):
+	def __init__(self, num_classes=10):
+		super().__init__()
+		model = models.resnet18()
+		model.fc = nn.Linear(512, 10)
+		self.model = model
+
+	def forward(self, x):
+		out = self.model(x)
+		return out
+
+class EffnetModel(nn.Module):
+	def __init__(self, num_classes=10) -> None:
+		super().__init__()
+		model = timm.create_model('efficientnet_b0', num_classes=10)
+		self.model = model
+
+	def forward(self, x):
+		out = self.model(x)
+		return out
+		

predict_image.py

@@ -0,0 +1,103 @@
+import cv2
+import torch
+from torchvision import transforms
+from model import ResnetModel, EffnetModel
+from face_module import get_face_coords
+from meter import Meter
+from utils import download_weights
+
+# statistics of imagenet dataset
+mean = [0.485, 0.456, 0.406]
+std = [0.229, 0.224, 0.225]
+
+# model wieghts url
+effb0_net_url = 'https://github.com/yMayanand/Emotion-Recognition/releases/download/v1.0.0/eff_b0.pt'
+res18_net_url = 'https://github.com/yMayanand/Emotion-Recognition/releases/download/v1.0.0/res18.pt'
+
+# transforms for image
+val_transform = transforms.Compose([
+		transforms.ToTensor(),
+        transforms.Resize((48, 48)),
+		transforms.Normalize(mean, std)
+	])
+
+def load_model(model_name):
+
+    # model for emotion classification
+    if model_name == 'effb0':
+        model = EffnetModel() 
+        fname = download_weights(effb0_net_url)
+    elif model_name == 'res18':
+        model = ResnetModel
+        fname = download_weights(res18_net_url)
+    else:
+        raise ValueError('Enter correct model_name')
+
+    # loading pretrained model
+    state_dict = torch.load(fname)
+    model.load_state_dict(state_dict['weights'])
+    return model
+
+# emotion classes
+emotions = ['neutral', 'happy :-)', 'surprise :-O', 'sad', 'angry >:(', 
+            "disgust D-':", 'fear', 'contempt', 'unknown', 'NF']
+
+# colors for text for each emotion classes
+colors = [(0, 128, 255), (255, 0, 255), (0, 255, 255), (255, 191, 0), (0, 0, 255),
+          (255, 255, 0), (0, 191, 255), (255, 0, 191), (255, 0, 191), (255, 0, 191)]
+
+def predict(image, save_path, model):
+    image = cv2.imread(image)
+    h, w, c = image.shape
+
+    # meter
+    m = Meter((w//2, h), w//5, (255, 0, 0))
+
+    # storing orignal image in bgr mode
+    orig_image = image
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    coords = get_face_coords(image)
+    if coords:
+        # getting bounding box coordinates for face
+        xmin, ymin, xmax, ymax = coords
+        model.eval()
+
+        image = image[ymin:ymax, xmin:xmax, :]
+
+        # check if face detected is not on edge of the screen
+        h, w, c = image.shape
+        if not (h and w):
+            idx = 9
+
+        image_tensor = val_transform(image).unsqueeze(0)
+        out = model(image_tensor)
+
+        # prediction emotion for detected face
+        pred = torch.argmax(out, dim=1)
+        idx = pred.item()
+        pred_emot = emotions[pred.item()]
+        color = colors[idx]
+
+        # drawing annotations on orignal bgr image
+        orig_image = cv2.rectangle(orig_image, (xmin, ymin), (xmax, ymax), (255, 0, 0), 1)
+    else:
+        idx = 9
+        pred_emot = 'Face Not Detected'
+        color = colors[-1]
+
+    orig_image = cv2.flip(orig_image, 1)
+    m.draw_meter(orig_image, idx)
+
+    cv2.imwrite(save_path, orig_image)
+    return orig_image
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--image_path', type=str, help='path to image location')
+    parser.add_argument('--model_name', type=str, default='effb0', help='name of the model')
+    parser.add_argument('--save_path', type=str, default='./result.jpg', help='path to save image')
+    args = parser.parse_args()
+    model = load_model(args.model_name)
+
+    predict(args.image_path, args.save_path, model)

requirements.txt

@@ -0,0 +1,10 @@
+###### Requirements without Version Specifiers ######
+
+gradio
+opencv-python
+mediapipe
+numpy
+torch 
+torchvision 
+torchaudio
+timm

utils.py

@@ -0,0 +1,31 @@
+import cv2
+import torch
+from urllib.request import urlretrieve
+
+def read_image(file):
+    """Reads the image file
+
+    Returns the numpy array.
+
+    Args:
+        file : path to the image
+
+    Returns:
+        (numpy.ndarray): image read as numpy array
+    """
+    image = cv2.imread(file)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    return image
+        
+def accuracy(predictions, ground_truth):
+    """Funtion to calculate accuracy of the model.
+    """
+    
+    _, preds = torch.max(predictions, dim=1)
+    score = (preds == ground_truth).float().mean()
+    return score.item()
+
+def download_weights(url):
+    fname = url.split('/')[-1]
+    urlretrieve(url, fname)
+    return fname