multipurpose-ai / app.py
Bipin Krishnan
added more description to the project
b6354a4
raw
history blame
3.74 kB
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 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)