Click the button to be redirected to the gradio app!
""" # sam = SAM() @app.get("/") def read_main(): return HTMLResponse(HTML) # @app.get("/foo") # def redirect(): # return RedirectResponse("/gradio") def detect_solar_panel(image) -> Results: # Perform inference results: Results = model(image) return results def segment_solar_panel(image) -> Results: # Perform inference results: Results = seg_model.predict(image, imgsz=(841, 595), retina_masks=True) return results def resize_and_pad(subject_image: Image.Image): # Resize subject image to 80% of 1200px while maintaining aspect ratio target_height = int(1200 * 0.8) aspect_ratio = subject_image.width / subject_image.height new_width = int(target_height * aspect_ratio) resized_subject = subject_image.resize((new_width, target_height), Image.LANCZOS) # Create a new transparent image new_image = Image.new("RGBA", (1200, 1200), (0, 0, 0, 0)) # Calculate the position to paste the resized subject image x = (1200 - new_width) // 2 y = (1200 - target_height) // 2 # Paste the resized subject image onto the transparent image new_image.paste(resized_subject, (x, y), resized_subject) # Save or return the PNG image png_image = new_image # Create a new image with a white background jpg_image = Image.new("RGB", (1200, 1200), (255, 255, 255)) jpg_image.paste(png_image, (0, 0), png_image) # Save or return the JPEG image return png_image, jpg_image def segment_image_core(img: np.ndarray | Image.Image) -> Image.Image: if type(img) is np.ndarray: img = Image.fromarray(img) results = segment_solar_panel(img) sections = [] for i, result in enumerate(results): print(f"Result {i}") result: Results try: h2, w2, c2 = result.orig_img.shape # Deal with boxes i = 0 for box in result.boxes: x1, y1, x2, y2 = box.xyxy[0].tolist() sections.append(((int(x1), int(y1), int(x2), int(y2)), f"{section_labels[0]} Bounding Box - index {i} - conf {box.conf}")) # Now the masks masks: Masks = result.masks try: mask = masks[i] cpu_mask = mask.cpu() squeezed_mask = cpu_mask.data.numpy() transposed_mask = squeezed_mask.transpose(1, 2, 0) kernel = cv2.getStructuringElement(cv2.MORPH_OPEN, (11, 11)) opened_mask = cv2.morphologyEx(transposed_mask, cv2.MORPH_OPEN, kernel, iterations=3) cv_mask = cv2.resize(opened_mask, (w2, h2)) image_mask = Image.fromarray((cv_mask * 255).astype(np.uint8)).filter(PIL.ImageFilter.GaussianBlur(1)) img_out = img.copy() img_out.putalpha(image_mask) img_out = img_out.crop((x1, y1, x2, y2)) png_img, jpg_img = resize_and_pad(img_out) sections.append((cv_mask, f"{section_labels[0]} Mask - Index: {i}")) except TypeError as e: print(f"Error processing image: {e}, probably no masks.") i += 1 except IndexError as e: print(f"Error processing image: {e}, probably no boxes.") return (img, sections), jpg_img def process_pdf_core(pdf) -> Image.Image: pdf = pdfium.PdfDocument(pdf) img_input.clear() # Get just the first page page = pdf[0] image = page.render(scale=4).to_pil() return image with gr.Blocks() as demo: section_labels = ['Solar Panel'] def segment_image(img): img_sections, jpg_img = segment_image_core(img) return img_sections # def process_image(img): # results = detect_solar_panel(img) # sections = [] # for result in results: # result: Results # # print(result) # try: # boxes = result.boxes.xyxy[0].tolist() # # Unpack boxes # x1, y1, x2, y2 = boxes # sections.append(((int(x1), int(y1), int(x2), int(y2)), f"{section_labels[0]} Bounding Box")) # #Create 4 centroids around the true centroid shifted by a delta value # delta = 0.3 # delta_x = (x2 - x1) * delta # delta_y = (y2 - y1) * delta # x_centroid = (x1 + x2) / 2 # y_centroid = (y1 + y2) / 2 # xtop_centroid = x_centroid # ytop_centroid = y_centroid + delta_y # xright_centroid = x_centroid + delta_x # yright_centroid = y_centroid # xbottom_centroid = x_centroid # ybottom_centroid = y_centroid - delta_y # xleft_centroid = x_centroid - delta_x # yleft_centroid = y_centroid # sam_mask, sam_scores = sam.segment(img, [[ # [xtop_centroid, ytop_centroid], # [xright_centroid, yright_centroid], # [xbottom_centroid, ybottom_centroid], # [xleft_centroid, yleft_centroid] # ]]) # squeezed_sam_mask_tensor = sam_mask[0].squeeze() # squeezed_sam_scores_tensor = sam_scores[0].squeeze() # print(f"sqeezed sam mask shape {squeezed_sam_mask_tensor.shape}") # print(f"sqeezed sam scores shape {squeezed_sam_scores_tensor.shape}") # for i in range(0, squeezed_sam_mask_tensor.shape[0]): # flat_mask = squeezed_sam_mask_tensor[i].numpy() # sections.append((flat_mask, f"{section_labels[0]} Mask {i} - Score: {squeezed_sam_scores_tensor[i]}")) # i += 1 # except IndexError as e: # print(f"Error processing image: {e}, probably no boxes.") # return (img, sections) def process_pdf(pdf): image = process_pdf_core(pdf) return segment_image(image) with gr.Row(): img_input = gr.Image(label="Upload Image", height=400) img_output = gr.AnnotatedImage(height=400) section_btn = gr.Button("Identify Solar Panel From Image") # Choose a random file in input directory gr.Examples( inputs = img_input, # examples = [os.path.join(image_dir, file) for file in random.sample(os.listdir(image_dir), 15)] examples = [os.path.join(image_dir, file) for file in os.listdir(image_dir)], ) with gr.Row(): pdf_input = gr.File(label="Upload PDF", file_types=['pdf'], height=200) pdf_btn = gr.Button("Identify Solar Panel from PDF") gr.Examples( inputs = pdf_input, examples = [os.path.join(pdf_dir, file) for file in os.listdir(pdf_dir)], ) section_btn.click(segment_image, [img_input], img_output) pdf_btn.click(process_pdf, [pdf_input], img_output) #Accept a PDF file, return a jpeg image @app.post("/uploadPdf", response_class=FileResponse) def extract_image(uploadFile: UploadFile) -> FileResponse: file = uploadFile.file.read() image = process_pdf_core(file) img_segments, jpeg_image = segment_image_core(image) id = str(uuid.uuid4()) filename = f"{cropped_dir}/cropped_{id}.jpg" jpeg_image.save(filename) return FileResponse(filename) app = gr.mount_gradio_app(app, demo, path="/gradio") if __name__ == "__main__": # app = gr.mount_gradio_app(app, demo, path="/gradio") uvicorn.run(app, port=7860) # demo.launch(share=True) # demo.launch(share=True, auth=(os.environ.get("GRADIO_USERNAME"), os.environ.get("GRADIO_PASSWORD")))