Update main.py

main.py

@@ -1,17 +1,10 @@
 import gradio as gr
-# import torch
-# import requests
-# from torchvision import transforms
 import cv2
 from geti_sdk.deployment import Deployment
 from geti_sdk.utils import show_image_with_annotation_scene
 
-# model = torch.hub.load("pytorch/vision:v0.6.0", "resnet18", pretrained=True).eval()
-# response = requests.get("https://git.io/JJkYN")
-# labels = response.text.split("\n")
 
-
-# Step 1: Load the deployment
+#Load the deployment
 deployment = Deployment.from_folder("deployment")
 deployment.load_inference_models(device="CPU")
 
@@ -31,27 +24,14 @@ def infer(image=None):
         return [None,'Error: No image provided']
 
     image = resize_image(image, 1200)
-    prediction = deployment.infer(image)
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    prediction = deployment.infer(image_rgb)
     output = show_image_with_annotation_scene(image, prediction, show_results=False)
     output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
     return [output, prediction.overview]
 
 
-# def predict(inp):
-#     inp = transforms.ToTensor()(inp).unsqueeze(0)
-#     with torch.no_grad():
-#         prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
-#         confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
-#     return confidences
-
-
 def run():
-    # demo = gr.Interface(
-    #     fn=predict,
-    #     inputs=gr.inputs.Image(type="pil"),
-    #     outputs=gr.outputs.Label(num_top_classes=3),
-    # )
-
     demo = gr.Interface(fn=infer, 
                         inputs=['image'],
                         outputs=['image', 'text'],