Spaces:

3martini
/

Solar-Eyes-Dockerized

Sleeping

App Files Files Community

3martini commited on Mar 15

Commit

68b5c0f

•

1 Parent(s): fd6f0a4

Upload folder using huggingface_hub

Browse files

Files changed (3) hide show

.gitignore +1 -0
app.py +37 -106
sandbox/sam_example.py +42 -0

.gitignore CHANGED Viewed

@@ -13,4 +13,5 @@ yolo*.pt
 **/.jpg
 pdf_downloads
 node_modules

 **/.jpg
 pdf_downloads
 node_modules
+.devcontainer/*.env

app.py CHANGED Viewed

@@ -3,7 +3,7 @@ import logging
 import json
 from fastapi import FastAPI, UploadFile
-from fastapi.responses import FileResponse, HTMLResponse, RedirectResponse
 import gradio as gr
 from PIL import Image
 import PIL
@@ -14,10 +14,15 @@ from ultralytics.engine.results import Results, Masks
 import uvicorn
 import cv2
 import uuid
-from functools import partial
 from openai import OpenAI
-PROMPT = "You are analyzing the spec sheet of a solar panel. Plese answer the following questions, format them as a JSON dictionary.\n"
 # from solareyes.sam import SAM
@@ -28,10 +33,6 @@ client = OpenAI(
 app = FastAPI()
-# Load the model
-# model: YOLO = YOLO('model/autodistill_best.pt')  # Path to trained model
-# seg_model: YOLO = YOLO('model/autodistill_best_seg.pt')  # Path to trained model
 # Directories
 image_dir = './pdf_images/'
 cropped_dir = './output/'
@@ -40,31 +41,6 @@ os.makedirs(image_dir, exist_ok=True)
 os.makedirs(cropped_dir, exist_ok=True)
 os.makedirs(pdf_dir, exist_ok=True)
-HTML = """
-<!DOCTYPE html>
-<html>
-<h1>Gradio Request Demo</h1>
-<p>Click the button to be redirected to the gradio app!</p>
-<button onclick="window.location.pathname='/gradio'">Redirect</button>
-</html>
-"""
-# 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 parse_pdf_text(file):
     pdf = pdfium.PdfDocument(file)
@@ -72,7 +48,9 @@ def parse_pdf_text(file):
     for page in pdf:
         textpage = page.get_textpage()
         text_all = textpage.get_text_bounded()
-        all_text += text_all
     #use openai to ask questions about text
     q1 = "What are module dimensions in L x W x H? Result key should be \"module_dimensions\""
@@ -101,6 +79,7 @@ def segment_solar_panel(image) -> Results:
     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)
@@ -168,14 +147,20 @@ def segment_image_core(img: np.ndarray | Image.Image) -> Image.Image:
     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:
@@ -185,93 +170,39 @@ with gr.Blocks() as demo:
     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)
-        img_output_pdf = gr.AnnotatedImage(height=400)
     with gr.Row():
         text_input = gr.Textbox(label="Enter Text", placeholder=PROMPT)
         text_output = gr.Textbox(label="Output", placeholder="Spec analysis will appear here")
-    pdf_btn = gr.Button("Identify Solar Panel from PDF")
     pdf_text_btn = gr.Button("Extract specs from PDF Text")
     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_pdf)
     pdf_text_btn.click(parse_pdf_text, [pdf_input], text_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"
@@ -282,10 +213,10 @@ def extract_image(uploadFile: UploadFile) -> FileResponse:
 #Accept a PDF file, return a text summary
 @app.post("/parsePdf")
 def parse_info(uploadFile: UploadFile):
     file = uploadFile.file.read()
-    logging.info(f"Received file {file}")
     answer = parse_pdf_text(file)
-    logging.info(f"Answer: {answer}")
     return {"answer": json.loads(answer)}

 import json
 from fastapi import FastAPI, UploadFile
+from fastapi.responses import FileResponse
 import gradio as gr
 from PIL import Image
 import PIL
 import uvicorn
 import cv2
 import uuid
 from openai import OpenAI
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG)
+PROMPT = """
+You are analyzing the spec sheet of a solar panel. If there is no text after the line \"PDF Extract Text Contents Below:\" report that there is no spec data provided as dictionary with a field called 'error'.
+If there is text, please answer the following questions, format them as a JSON dictionary. Include the units of dimensions, weight, and cable lengths.\n
+"""
 # from solareyes.sam import SAM
 app = FastAPI()
 # Directories
 image_dir = './pdf_images/'
 cropped_dir = './output/'
 os.makedirs(cropped_dir, exist_ok=True)
 os.makedirs(pdf_dir, exist_ok=True)
 def parse_pdf_text(file):
     pdf = pdfium.PdfDocument(file)
     for page in pdf:
         textpage = page.get_textpage()
         text_all = textpage.get_text_bounded()
+        all_text += text_all
+    logger.debug(f"Text: {all_text}")
     #use openai to ask questions about text
     q1 = "What are module dimensions in L x W x H? Result key should be \"module_dimensions\""
     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)
     return (img, sections), jpg_img
+def pdf_to_image(pdf, end, start = 0) -> list[Image.Image]:
     pdf = pdfium.PdfDocument(pdf)
+    page_images = []
+    if not end:
+        end = len(pdf)  # get the number of pages in the document
+    for i in range(start, end):
+        page = pdf[i]
+        page_image = page.render(scale=4).to_pil()
+        page_images.append(page_image)
+    return page_images
+def pdf_first_page_to_image(pdf) -> Image.Image:
+    return pdf_to_image(pdf, 1, 0)[0]
 with gr.Blocks() as demo:
     def segment_image(img):
         img_sections, jpg_img = segment_image_core(img)
         return img_sections
     def process_pdf(pdf):
+        image = pdf_first_page_to_image(pdf)
         return segment_image(image)
+    pdf_input = gr.File(label="Upload PDF", file_types=['pdf'], height=100)
+    pdf_image = gr.Gallery(label="PDF Page Images")
+    pdf_to_image_btn = gr.Button("Convert PDF to Image")
     with gr.Row():
+        img_output_pdf = gr.AnnotatedImage(label="Extracted product image", height=400)
+    pdf_extract_btn = gr.Button("Identify Solar Panel from PDF")
     with gr.Row():
         text_input = gr.Textbox(label="Enter Text", placeholder=PROMPT)
         text_output = gr.Textbox(label="Output", placeholder="Spec analysis will appear here")
     pdf_text_btn = gr.Button("Extract specs from PDF Text")
     gr.Examples(
         inputs = pdf_input,
         examples = [os.path.join(pdf_dir, file) for file in os.listdir(pdf_dir)],
     )
+    pdf_extract_btn.click(process_pdf, [pdf_input], img_output_pdf)
     pdf_text_btn.click(parse_pdf_text, [pdf_input], text_output)
+    pdf_to_image_btn.click(pdf_to_image, [pdf_input], pdf_image)
 #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 = pdf_to_image(file)
     img_segments, jpeg_image = segment_image_core(image)
     id = str(uuid.uuid4())
     filename = f"{cropped_dir}/cropped_{id}.jpg"
 #Accept a PDF file, return a text summary
 @app.post("/parsePdf")
 def parse_info(uploadFile: UploadFile):
+    logger.info(f"Receiving file {uploadFile.filename}")
     file = uploadFile.file.read()
+    logger.info(f"Received file {uploadFile.filename}")
     answer = parse_pdf_text(file)
     return {"answer": json.loads(answer)}

sandbox/sam_example.py ADDED Viewed

	@@ -0,0 +1,42 @@

+    # 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)