Spaces:
Build error
Build error
| import os | |
| from langchain.llms import OpenAI, OpenAIChat | |
| os.system("pip install -U gradio") | |
| import sys | |
| import gradio as gr | |
| os.system( | |
| "pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html" | |
| ) | |
| # clone and install Detic | |
| os.system( | |
| "git clone https://github.com/facebookresearch/Detic.git --recurse-submodules" | |
| ) | |
| os.chdir("Detic") | |
| # Install detectron2 | |
| import torch | |
| # Some basic setup: | |
| # Setup detectron2 logger | |
| import detectron2 | |
| from detectron2.utils.logger import setup_logger | |
| setup_logger() | |
| # import some common libraries | |
| import sys | |
| import numpy as np | |
| import os, json, cv2, random | |
| # import some common detectron2 utilities | |
| from detectron2 import model_zoo | |
| from detectron2.engine import DefaultPredictor | |
| from detectron2.config import get_cfg | |
| from detectron2.utils.visualizer import Visualizer | |
| from detectron2.data import MetadataCatalog, DatasetCatalog | |
| # Detic libraries | |
| sys.path.insert(0, "third_party/CenterNet2/projects/CenterNet2/") | |
| sys.path.insert(0, "third_party/CenterNet2/") | |
| from centernet.config import add_centernet_config | |
| from detic.config import add_detic_config | |
| from detic.modeling.utils import reset_cls_test | |
| from PIL import Image | |
| # Build the detector and download our pretrained weights | |
| cfg = get_cfg() | |
| add_centernet_config(cfg) | |
| add_detic_config(cfg) | |
| cfg.MODEL.DEVICE = "cpu" | |
| cfg.merge_from_file("configs/Detic_LCOCOI21k_CLIP_SwinB_896b32_4x_ft4x_max-size.yaml") | |
| cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detic/Detic_LCOCOI21k_CLIP_SwinB_896b32_4x_ft4x_max-size.pth" | |
| cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model | |
| cfg.MODEL.ROI_BOX_HEAD.ZEROSHOT_WEIGHT_PATH = "rand" | |
| cfg.MODEL.ROI_HEADS.ONE_CLASS_PER_PROPOSAL = ( | |
| True # For better visualization purpose. Set to False for all classes. | |
| ) | |
| predictor = DefaultPredictor(cfg) | |
| BUILDIN_CLASSIFIER = { | |
| "lvis": "datasets/metadata/lvis_v1_clip_a+cname.npy", | |
| "objects365": "datasets/metadata/o365_clip_a+cnamefix.npy", | |
| "openimages": "datasets/metadata/oid_clip_a+cname.npy", | |
| "coco": "datasets/metadata/coco_clip_a+cname.npy", | |
| } | |
| BUILDIN_METADATA_PATH = { | |
| "lvis": "lvis_v1_val", | |
| "objects365": "objects365_v2_val", | |
| "openimages": "oid_val_expanded", | |
| "coco": "coco_2017_val", | |
| } | |
| session_token = os.environ.get("SessionToken") | |
| def generate_caption(object_list_str, api_key, temperature): | |
| query = f"You are an intelligent image captioner. I will hand you the objects and their position, and you should give me a detailed description for the photo. In this photo we have the following objects\n{object_list_str}" | |
| llm = OpenAIChat( | |
| model_name="gpt-3.5-turbo", openai_api_key=api_key, temperature=temperature | |
| ) | |
| try: | |
| caption = llm(query) | |
| caption = caption.strip() | |
| except: | |
| caption = "Sorry, something went wrong!" | |
| return caption | |
| def inference(img, vocabulary, api_key, temperature): | |
| metadata = MetadataCatalog.get(BUILDIN_METADATA_PATH[vocabulary]) | |
| classifier = BUILDIN_CLASSIFIER[vocabulary] | |
| num_classes = len(metadata.thing_classes) | |
| reset_cls_test(predictor.model, classifier, num_classes) | |
| im = cv2.imread(img) | |
| outputs = predictor(im) | |
| v = Visualizer(im[:, :, ::-1], metadata) | |
| out = v.draw_instance_predictions(outputs["instances"].to("cpu")) | |
| detected_objects = [] | |
| object_list_str = [] | |
| box_locations = outputs["instances"].pred_boxes | |
| box_loc_screen = box_locations.tensor.cpu().numpy() | |
| for i, box_coord in enumerate(box_loc_screen): | |
| x0, y0, x1, y1 = box_coord | |
| width = x1 - x0 | |
| height = y1 - y0 | |
| predicted_label = metadata.thing_classes[outputs["instances"].pred_classes[i]] | |
| detected_objects.append( | |
| { | |
| "prediction": predicted_label, | |
| "x": int(x0), | |
| "y": int(y0), | |
| "w": int(width), | |
| "h": int(height), | |
| } | |
| ) | |
| object_list_str.append( | |
| f"{predicted_label} - X:({int(x0)} Y: {int(y0)} Width {int(width)} Height: {int(height)})" | |
| ) | |
| if api_key is not None: | |
| gpt_response = generate_caption(object_list_str, api_key, temperature) | |
| else: | |
| gpt_response = "Please paste your OpenAI key to use" | |
| return ( | |
| Image.fromarray(np.uint8(out.get_image())).convert("RGB"), | |
| gpt_response, | |
| ) | |
| with gr.Blocks() as demo: | |
| with gr.Column(): | |
| gr.Markdown("# Image Captioning using Detic and ChatGPT with LangChain 🦜️🔗") | |
| gr.Markdown( | |
| "Use Detic to detect objects in an image and then use `gpt-3.5-turbo` to describe the image." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| inp = gr.Image(label="Input Image", type="filepath") | |
| with gr.Column(): | |
| openai_api_key_textbox = gr.Textbox( | |
| placeholder="Paste your OpenAI API key (sk-...)", | |
| show_label=False, | |
| lines=1, | |
| type="password", | |
| ) | |
| temperature = gr.Slider(0, 1, 0.1, label="Temperature") | |
| vocab = gr.Dropdown( | |
| ["lvis", "objects365", "openimages", "coco"], | |
| label="Detic Vocabulary", | |
| value="lvis", | |
| ) | |
| btn_detic = gr.Button("Run Detic and ChatGPT") | |
| with gr.Column(): | |
| output_desc = gr.Textbox(label="Description Description", lines=5) | |
| outviz = gr.Image(label="Visualization", type="pil") | |
| btn_detic.click( | |
| fn=inference, | |
| inputs=[inp, vocab, openai_api_key_textbox, temperature], | |
| outputs=[outviz, output_desc], | |
| ) | |
| demo.launch(debug=False) | |