Update app.py

app.py

@@ -1,6 +1,4 @@
 import os
-# os.system('wget https://huggingface.co/spaces/An-619/FastSAM/resolve/main/weights/FastSAM.pt')
-
 import yolov5
 
 # load model
@@ -41,7 +39,7 @@ def read_license_number(img):
                for bbox in boxes]
 
 from transformers import CLIPProcessor, CLIPModel
-model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+vit_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
 processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
 def zero_shot_classification(image, labels):
@@ -49,7 +47,7 @@ def zero_shot_classification(image, labels):
                        images=image,
                        return_tensors="pt",
                        padding=True)
-    outputs = model(**inputs)
+    outputs = vit_model(**inputs)
     logits_per_image = outputs.logits_per_image # this is the image-text similarity score
     return logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
 
@@ -75,33 +73,34 @@ def check_solarplant_broken(image):
     idx = probs.argmax().item()
     return zero_shot_class_labels[idx].split(" ")[1-idx]
 
-# from fastsam import FastSAM, FastSAMPrompt
+from fastsam import FastSAM, FastSAMPrompt
+os.system('wget https://huggingface.co/spaces/An-619/FastSAM/resolve/main/weights/FastSAM.pt')
 
-# model = FastSAM('./FastSAM.pt')
-# DEVICE = 'cpu'
-# def segment_solar_panel(img):
-#     # os.system('python Inference.py --model_path FastSAM.pt --img_path bus.jpg  --text_prompt "solar panel grids"')
-#     img = img.convert("RGB")
+model = FastSAM('./FastSAM.pt')
+DEVICE = 'cpu'
+def segment_solar_panel(img):
+    # os.system('python Inference.py --model_path FastSAM.pt --img_path bus.jpg  --text_prompt "solar panel grids"')
+    img = img.convert("RGB")
 
-#     everything_results = model(img, device=DEVICE, retina_masks=True, imgsz=1024, conf=0.4, iou=0.9,)
-#     prompt_process = FastSAMPrompt(img, everything_results, device=DEVICE)
+    everything_results = model(img, device=DEVICE, retina_masks=True, imgsz=1024, conf=0.4, iou=0.9,)
+    prompt_process = FastSAMPrompt(img, everything_results, device=DEVICE)
     
-#     # everything prompt
-#     ann = prompt_process.everything_prompt()
+    # everything prompt
+    ann = prompt_process.everything_prompt()
     
-#     # bbox default shape [0,0,0,0] -> [x1,y1,x2,y2]
-#     ann = prompt_process.box_prompt(bbox=[[200, 200, 300, 300]])
+    # bbox default shape [0,0,0,0] -> [x1,y1,x2,y2]
+    ann = prompt_process.box_prompt(bbox=[[200, 200, 300, 300]])
     
-#     # text prompt
-#     ann = prompt_process.text_prompt(text='solar panel grids')
+    # text prompt
+    ann = prompt_process.text_prompt(text='solar panel grids')
     
-#     # point prompt
-#     # points default [[0,0]] [[x1,y1],[x2,y2]]
-#     # point_label default [0] [1,0] 0:background, 1:foreground
-#     ann = prompt_process.point_prompt(points=[[620, 360]], pointlabel=[1])
+    # point prompt
+    # points default [[0,0]] [[x1,y1],[x2,y2]]
+    # point_label default [0] [1,0] 0:background, 1:foreground
+    ann = prompt_process.point_prompt(points=[[620, 360]], pointlabel=[1])
     
-#     prompt_process.plot(annotations=ann,output_path='./bus.jpg',)
-#     return Image.Open('./bus.jpg')
+    prompt_process.plot(annotations=ann,output_path='./bus.jpg',)
+    return Image.Open('./bus.jpg')
 
 
 import gradio as gr
@@ -109,8 +108,7 @@ import gradio as gr
 def greet(img):
     lns = read_license_number(img)
     if len(lns):
-        # seg = segment_solar_panel(img)
-        seg = img
+        seg = segment_solar_panel(img)
         return (seg,
                 "車牌： " + '; '.join(lns) + "\n\n" \
                 + "類型： "+ check_solarplant_installed_by_image(img, True) + "\n\n" \