add project code

app.py

@@ -0,0 +1,201 @@
+import base64
+import concurrent.futures
+import os
+from io import BytesIO
+
+import cv2
+import gradio as gr
+import numpy as np
+import requests
+import supervision as sv
+from inference_sdk import InferenceHTTPClient
+from openai import OpenAI
+
+CLIENT = InferenceHTTPClient(
+    api_url="http://detect.roboflow.com",
+    api_key=os.environ["ROBOFLOW_API_KEY"],
+)
+
+openai_client = OpenAI()
+
+
+def process_mask(region, task_id):
+    region = cv2.rotate(region, cv2.ROTATE_90_CLOCKWISE)
+
+    cv2.imwrite(f"region_{task_id}.jpg", region)
+
+    # change channels
+    region = cv2.cvtColor(region, cv2.COLOR_BGR2RGB)
+
+    base64_image = base64.b64encode(
+        BytesIO(cv2.imencode(".jpg", region)[1]).read()
+    ).decode("utf-8")
+
+    response = openai_client.chat.completions.create(
+        model="gpt-4-vision-preview",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "Read the text on the book spine. Only say the book cover title and author if you can find them. Say the book that is most prominent. Return the format [title] [author], with no punctuation.",
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
+                    },
+                ],
+            }
+        ],
+        max_tokens=300,
+    )
+
+    print(response.choices[0].message.content.rstrip("Title:").replace("\n", " "))
+
+    return response.choices[0].message.content.rstrip("Title:").replace("\n", " ")
+
+
+def process_book_with_google_books(book):
+    response = requests.get(
+        f"https://www.googleapis.com/books/v1/volumes?q={book}",
+        headers={"User-Agent": "Mozilla/5.0"},
+    )
+    response = response.json()
+
+    isbn, author, link = "NULL", "NULL", "NULL"
+
+    try:
+        isbn = response["items"][0]["volumeInfo"]["industryIdentifiers"][0][
+            "identifier"
+        ]
+        if (
+            "volumeInfo" in response["items"][0]
+            and "authors" in response["items"][0]["volumeInfo"]
+        ):
+            author = response["items"][0]["volumeInfo"]["authors"][0]
+        link = response["items"][0]["volumeInfo"]["infoLink"]
+    except:
+        pass
+
+    return isbn, author, link
+
+
+# define function that accepts an image
+def detect_books(image):
+    # infer on a local image
+    results = CLIENT.infer(image, model_id="open-shelves/4")
+    results = sv.Detections.from_inference(results)
+
+    mask_annotator = sv.MaskAnnotator()
+    annotated_image = mask_annotator.annotate(scene=image, detections=results)
+
+    masks_isolated = []
+
+    masks_to_xyxys = sv.mask_to_xyxy(masks=results.mask)
+
+    polygons = [sv.mask_to_polygons(mask) for mask in results.mask]
+
+    for mask in results.mask:
+        masked_region = np.zeros_like(image)
+        masked_region[mask] = image[mask]
+        masks_isolated.append(masked_region)
+
+    print("Calculated masks...")
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        tasks = [
+            executor.submit(process_mask, region, task_id)
+            for task_id, region in enumerate(masks_isolated)
+        ]
+        books = [task.result() for task in tasks]
+
+    print("Processed books...")
+
+    links = []
+
+    isbns = []
+    authors = []
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        tasks = [
+            executor.submit(process_book_with_google_books, book) for book in books
+        ]
+
+        for task in tasks:
+            isbn, author, link = task.result()
+            isbns.append(isbn)
+            authors.append(author)
+            links.append(link)
+
+    print("Processed books with Google Books...")
+
+    annotations = [
+        {
+            "title": title,
+            "author": author,
+            "isbn": isbn,
+            "polygons": [polygon.tolist() for polygon in polygon_list],
+            "xyxy": xyxy.tolist(),
+            "link": link,
+        }
+        for title, author, isbn, polygon_list, xyxy, link in zip(
+            books, authors, isbns, polygons, results.xyxy, links
+        )
+        if "sorry" not in title.lower() and "NULL" not in title
+    ]
+
+    width, height = image.shape[1], image.shape[0]
+
+    svg = f"""<div class="image-container"><img src="image.jpeg" height="{height}" width="{width}">
+        
+        <svg width="{width}" height="{height}">"""
+
+    for annotation in annotations:
+        polygons = annotation["polygons"][0]
+        svg += f"""<polygon points="{', '.join([f'{x},{y}' for x, y in polygons])}" fill="transparent" stroke="red" stroke-width="2"
+        onclick="window.location.href='{annotation['link']}';"></polygon>"""
+
+    svg += "</svg>"
+    svg += """
+        <style>
+        .image-container {
+            position: relative;
+            height: HEIGHTpx;
+            width: WIDTHpx;
+        }
+        .image-container img, .image-container svg {
+            position: absolute;
+            left: 0;
+            top: 0;
+            width: 100%;
+            height: auto;
+        }
+        .image-container svg {
+            z-index: 1;
+        }
+        </style></div>""".replace(
+        "HEIGHT", str(height)
+    ).replace(
+        "WIDTH", str(width)
+    )
+
+    return annotated_image, books, isbns, svg
+
+
+iface = gr.Interface(
+    fn=detect_books,
+    description="Use Open Shelves to detect books in an image. The model will return the annotated image with the detected books, the titles of the books, and the ISBNs of the books.",
+    inputs=gr.components.Image(label="Input Image"),
+    # outputs should be an image and a list of text
+    outputs=[
+        gr.components.Image(label="Annotated Image"),
+        gr.components.Textbox(label="Detected Books"),
+        gr.components.Textbox(label="ISBNs"),
+        gr.components.Textbox(label="SVG"),
+    ],
+    title="Open Shelves",
+    allow_flagging=False,
+    theme="huggingface",
+)
+iface.launch()

requirements.txt

@@ -0,0 +1,6 @@
+gradio
+requests
+numpy
+supervision
+inference-sdk
+openai