Spaces:
Runtime error
Runtime error
| import os | |
| # import yolov5 | |
| # # load model | |
| # model = yolov5.load('keremberke/yolov5m-license-plate') | |
| # # set model parameters | |
| # model.conf = 0.5 # NMS confidence threshold | |
| # model.iou = 0.25 # NMS IoU threshold | |
| # model.agnostic = False # NMS class-agnostic | |
| # model.multi_label = False # NMS multiple labels per box | |
| # model.max_det = 1000 # maximum number of detections per image | |
| # # set image | |
| # def license_plate_detect(img): | |
| # # perform inference | |
| # results = model(img, size=640) | |
| # # inference with test time augmentation | |
| # results = model(img, augment=True) | |
| # # parse results | |
| # if len(results.pred): | |
| # predictions = results.pred[0] | |
| # boxes = predictions[:, :4] # x1, y1, x2, y2 | |
| # scores = predictions[:, 4] | |
| # categories = predictions[:, 5] | |
| # return boxes | |
| # from PIL import Image | |
| # # image = Image.open(img) | |
| # import pytesseract | |
| # def read_license_number(img): | |
| # boxes = license_plate_detect(img) | |
| # if boxes: | |
| # return [pytesseract.image_to_string( | |
| # image.crop(bbox.tolist())) | |
| # for bbox in boxes] | |
| from transformers import CLIPProcessor, CLIPModel | |
| 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): | |
| inputs = processor(text=labels, | |
| images=image, | |
| return_tensors="pt", | |
| padding=True) | |
| 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 | |
| # installed_list = [] | |
| # # image = Image.open(requests.get(url, stream=True).raw) | |
| # def check_solarplant_installed_by_license(license_number_list): | |
| # if len(installed_list): | |
| # return [license_number in installed_list | |
| # for license_number in license_number_list] | |
| def check_solarplant_installed_by_image(image, output_label=False): | |
| zero_shot_class_labels = ["bus with solar panel grids", | |
| "bus without solar panel grids"] | |
| probs = zero_shot_classification(image, zero_shot_class_labels) | |
| if output_label: | |
| return zero_shot_class_labels[probs.argmax().item()] | |
| return probs.argmax().item() == 0 | |
| # def check_solarplant_broken(image): | |
| # zero_shot_class_labels = ["white broken solar panel", | |
| # "normal black solar panel grids"] | |
| # probs = zero_shot_classification(image, zero_shot_class_labels) | |
| # idx = probs.argmax().item() | |
| # return zero_shot_class_labels[idx].split(" ")[1-idx] | |
| 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") | |
| 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() | |
| # 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') | |
| # 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') | |
| import gradio as gr | |
| def greet(img): | |
| if check_solarplant_installed_by_image(img): | |
| seg = segment_solar_panel(img) | |
| return (seg, 'ε試εε²ε€ͺι½θ½ζΏι¨ε') | |
| # return (seg, | |
| # "θ»ηοΌ " + '; '.join(lns) + "\n\n" \ | |
| # + "ι‘εοΌ "+ check_solarplant_installed_by_image(img, True) + "\n\n" \ | |
| # + "ηζ οΌ" + check_solarplant_broken(img)) | |
| return (img, "ζ²ζε€ͺι½θ½ζΏι¨εεε²") | |
| iface = gr.Interface(fn=greet, inputs="image", outputs=["image", "text"]) | |
| iface.launch() |