|
import os |
|
os.system('wget https://huggingface.co/spaces/An-619/FastSAM/resolve/main/weights/FastSAM.pt') |
|
|
|
import yolov5 |
|
|
|
|
|
model = yolov5.load('keremberke/yolov5m-license-plate') |
|
|
|
|
|
model.conf = 0.5 |
|
model.iou = 0.25 |
|
model.agnostic = False |
|
model.multi_label = False |
|
model.max_det = 1000 |
|
|
|
|
|
def license_plate_detect(img): |
|
|
|
results = model(img, size=640) |
|
|
|
|
|
results = model(img, augment=True) |
|
|
|
|
|
if len(results.pred): |
|
predictions = results.pred[0] |
|
boxes = predictions[:, :4] |
|
scores = predictions[:, 4] |
|
categories = predictions[:, 5] |
|
return boxes |
|
|
|
from PIL import Image |
|
|
|
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 |
|
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 = model(**inputs) |
|
logits_per_image = outputs.logits_per_image |
|
return logits_per_image.softmax(dim=1) |
|
|
|
installed_list = [] |
|
|
|
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][1-idx] |
|
|
|
from fastsam import FastSAM, FastSAMPrompt |
|
|
|
model = FastSAM('./FastSAM.pt') |
|
DEVICE = 'cpu' |
|
def segment_solar_panel(img): |
|
|
|
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) |
|
|
|
|
|
ann = prompt_process.everything_prompt() |
|
|
|
|
|
ann = prompt_process.box_prompt(bbox=[[200, 200, 300, 300]]) |
|
|
|
|
|
ann = prompt_process.text_prompt(text='solar panel grids') |
|
|
|
|
|
|
|
|
|
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): |
|
lns = read_license_number(img) |
|
if len(lns): |
|
seg = segment_solar_panel(img) |
|
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() |