multipurpose-ai / app.py
bipin's picture
added a clarification to the description
6d2884b
import gradio as gr
from transformers import pipeline
from glob import glob
from PIL import Image
import os
from icrawler.builtin import GoogleImageCrawler
def download_image(query, out_file):
google_crawler = GoogleImageCrawler(storage={'root_dir': './'})
google_crawler.crawl(keyword=query, max_num=1, overwrite=True)
os.rename(glob("./000001.*")[0], out_file)
def generate_story(prompt):
story = storygen(f"{prompt}")[0]['generated_text']
return story
def start_neural_style_transfer(img1_name, img2_name):
img1_filename, img2_filename = "style_transfer_1.jpg", "style_transfer_2.jpg"
download_image(img1_name, img1_filename)
download_image(img2_name, img2_filename)
styled_image = nst(img1_filename, img2_filename)
pil_img = Image.open(styled_image)
return pil_img
def detect_objects(file_name):
out_img = detectron(file_name)
pil_img = Image.open(out_img)
return pil_img
def main(text_input):
text_output, image_output, metadata = None, None, None
task_type_q = f"User: {text_input}\nWhat is the task the user is asking to do?\n \
- story generation task\n \
- image style transfer task\n \
- object detection task"
task = t0pp(task_type_q)
task = task.lower().replace('.', '')
if task=="story generation task":
story_prompt = t0pp(f"User: {text_input}\nWhat story is the user asking for?")
text_output = generate_story(story_prompt)
metadata = f"Prompt used to generate the story:\n{story_prompt}"
elif task=="image style transfer task":
img1_name = t0pp(f"User: {text_input}\nWhat is the name of the picture to which style is to be tranferred?")
img2_name = t0pp(f"User: {text_input}\nWhat is the name of the picture from which style is to be tranferred?")
image_output = start_neural_style_transfer(img1_name, img2_name)
metadata = f"Image from which style is to be transferred: {img2_name}\nImage to which style is to be transferred: {img1_name}"
elif task=="object detection task":
img_file = "object_detection.jpg"
img_name = t0pp(f"User: {text_input}\nWhat image is the user referring to?")
download_image(img_name, img_file)
image_output = detect_objects(img_file)
metadata = f"Image from which objects are to be detected: {img_name}"
return text_output, image_output
if __name__=="__main__":
t0pp = gr.Interface.load("huggingface/bigscience/T0pp")
storygen = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
nst = gr.Interface.load("spaces/luca-martial/neural-style-transfer")
detectron = gr.Interface.load("spaces/akhaliq/Detectron2")
title = "Multipurpose AI"
description = "With this demo you can do image style transfer, story generation and object detection(as of now, detects all the possible objects in the image) using a single text box. Use the examples provided below to get started."
article = "[Bigscience T0pp](https://huggingface.co/bigscience/T0pp) | [story generator](https://huggingface.co/pranavpsv/genre-story-generator-v2) | [style transfer](https://huggingface.co/spaces/luca-martial/neural-style-transfer) | [object detection](https://huggingface.co/spaces/akhaliq/Detectron2)"
examples = [["Find an image of taj mahal and transfer the style of starry night to it"],
["Find an image of a football match and detect the objects in that image"],
["Give me a story about a hard working farmer"]]
gr.Interface(
main,
title=title,
description=description,
article=article,
inputs=gr.inputs.Textbox(lines=5, label="Input"),
outputs=[gr.outputs.Textbox(label="Output"), gr.outputs.Image(label="Ouptut"),],
examples=examples,
enable_queue=True,
).launch(debug=True)