fixes

app.py

@@ -12,6 +12,9 @@ from mediapipe import solutions
 from PIL import Image
 
 import torch, torchvision
+import torchvision.transforms as T
+from huggingface_hub import hf_hub_download
+
 
 import matplotlib
 matplotlib.use("Agg")
@@ -19,10 +22,6 @@ import matplotlib.pyplot as plt
 
 cropped_image = []
 analyzed_image = []
-
-# colors for visualization
-COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
-          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
      
 finetuned_classes = [
       'iris',
@@ -41,20 +40,17 @@ options = vision.FaceLandmarkerOptions(base_options=base_options,
                                        num_faces=1)
 detector = vision.FaceLandmarker.create_from_options(options)
 
-model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True)
-model.eval();
+# Loading the model 
+model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=False, num_classes=1)
+hf_hub_download(repo_id="zivpollak/ECXV001", filename="checkpoint.pth", local_dir='.')
+checkpoint = torch.load('checkpoint.pth', map_location='cpu') 
+model.load_state_dict(checkpoint['model'], strict=False) 
+model.eval()
 
 def video_identity(video):
     return video
 
 
-#demo = gr.Interface(video_identity,
-#                    gr.Video(shape = (1000,1000), source="webcam"),
-#                    "playable_video")
-
-
-import torchvision.transforms as T
-
 # standard PyTorch mean-std input image normalization
 transform = T.Compose([
     T.Resize(800),
@@ -65,11 +61,12 @@ transform = T.Compose([
 def handle_image(input_image):
     global cropped_image, analyzed_image
     cv2.imwrite("image.jpg", input_image) 
-    image = mp.Image.create_from_file("image.jpg")
+    #image = mp.Image.create_from_file("image.jpg")
+    image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(input_image))
 
-    detection_result = detector.detect(image)
     cropped_image = image.numpy_view().copy()
     analyzed_image = image.numpy_view().copy()
+    detection_result = detector.detect(image)
 
     face_landmarks_list = detection_result.face_landmarks
 
@@ -87,26 +84,9 @@ def handle_image(input_image):
     cv2.circle(input_image, (p2[0], p2[1]), 10, (0, 0, 255), -1)
     cropped_image = cropped_image[p1[1]:p2[1], p1[0]:p2[0]]
     
-    run_worflow(cropped_image, model)
-    
-    return (cropped_image)
-
-def load_model():
-    print('load model')
-    '''
-    model = torch.hub.load('facebookresearch/detr',
-                       'detr_resnet50',
-                       pretrained=False,
-                       num_classes=1)
-
-    checkpoint = torch.load('outputs/checkpoint.pth',
-                            map_location='cpu')
+    output_image = run_worflow(cropped_image, model)
+    return (output_image)
 
-    model.load_state_dict(checkpoint['model'],
-                      strict=False)
-
-    model.eval();
-    '''
 def filter_bboxes_from_outputs(img,
                                outputs,
                                threshold=0.7
@@ -124,22 +104,13 @@ def filter_bboxes_from_outputs(img,
   return probas_to_keep, bboxes_scaled
 
 
-def plot_finetuned_results(pil_img, prob=None, boxes=None):
-    plt.figure(figsize=(16,10))
-    plt.imshow(pil_img)
-    ax = plt.gca()
-    colors = COLORS * 100
+def plot_finetuned_results(img, prob=None, boxes=None):
     if prob is not None and boxes is not None:
-      for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
-          ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
-                                    fill=False, color=c, linewidth=3))
-          cl = p.argmax()
-          #text = f'{finetuned_classes[cl]}: {p[cl]:0.2f}'
-          text = 'results'
-          ax.text(xmin, ymin, text, fontsize=15,
-                  bbox=dict(facecolor='yellow', alpha=0.5))
-    plt.axis('off')
-    plt.show()
+      for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
+          print("adding rectangle")
+          cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 255, 255), 1)
+    return img      
+
      
 def rescale_bboxes(out_bbox, size):
     print (size)
@@ -167,15 +138,20 @@ def run_worflow(my_image, my_model):
   # propagate through the model
   outputs = my_model(img)
 
-  for threshold in [0.2, 0.2]:
+  output_image =   cv2.imread("img1.jpg")
+
+  for threshold in [0.4, 0.4]:
     
     probas_to_keep, bboxes_scaled = filter_bboxes_from_outputs(my_image,
                                                                outputs,
                                                               threshold=threshold)
 
-    plot_finetuned_results(my_image,
+    print(bboxes_scaled)
+    output_image = plot_finetuned_results(output_image,
                            probas_to_keep, 
                            bboxes_scaled)
+    
+    return output_image