Spaces:
Sleeping
Sleeping
import gradio as gr | |
import torch | |
import torchvision.transforms as transforms | |
from sentence_transformers import SentenceTransformer, util | |
import json | |
import os | |
import matplotlib.pyplot as plt | |
import random | |
import torch | |
import torch.nn as nn | |
from generatorModel import Generator | |
noise_dim = 16 | |
embed_dim = 384 | |
embed_out_dim = 256 | |
device = 'cpu' | |
generator = Generator(channels=3, embed_dim=embed_dim, noise_dim=noise_dim, embed_out_dim=embed_out_dim).to(device) | |
# Path to your .pth file | |
gen_weight = 'generator_20240421_3.pth' | |
# Load the weights | |
weights_gen = torch.load(gen_weight, map_location=torch.device(device)) | |
# Apply the weights to your model | |
generator.load_state_dict(weights_gen) | |
model = SentenceTransformer('sentence-transformers/all-MiniLM-L12-v2') | |
def load_embedding(model): | |
# Load your model and other components here | |
with open(os.path.join("descriptions.json"), 'r') as file: | |
dataset = json.load(file) | |
classes = [e["text"] for e in dataset] | |
embeddings_list = {cls: model.encode(cls, convert_to_tensor=True) for cls in classes} | |
return embeddings_list | |
def generate_image(caption): | |
embeddings = load_embedding(model) | |
noise_dim = 16 | |
results = [(util.pytorch_cos_sim(model.encode(caption, convert_to_tensor=True), embeddings[cls]).item(), cls) for cls in classes] | |
sorted_results = sorted(results, key=lambda x: x[0], reverse=True)[:5] | |
threshold = 0.40 | |
coeff = 0.89 | |
if sorted_results[0][0] <= threshold: | |
caption = sorted_results[0][1] | |
results = [(util.pytorch_cos_sim(model.encode(caption, convert_to_tensor=True), embeddings[cls]).item(), cls) for cls in classes] | |
sorted_results = sorted(results, key=lambda x: x[0], reverse=True)[:5] | |
if sorted_results[0][0] >= 0.99: | |
coeff = 0.85 | |
last_score = sorted_results[0][0] | |
filtered_results = [] | |
for score, cls in sorted_results: | |
if score >= last_score * coeff: | |
filtered_results.append((score, cls)) | |
last_score = score | |
else: | |
break | |
items = [cls for score, cls in filtered_results] | |
probabilities = [score for score, cls in filtered_results] | |
sampled_item = random.choices(items, weights=probabilities, k=1)[0] | |
noise = torch.randn(1, noise_dim, 1, 1, device=device) # Adjust noise_dim if different | |
fake_images = generator(noise, embeddings[sampled_item].unsqueeze(0).unsqueeze(0)) | |
img = fake_images.squeeze(0).permute(1, 2, 0).cpu().detach().numpy() | |
img = (img - img.min()) / (img.max() - img.min()) | |
return img | |
iface = gr.Interface(fn=generate_image, | |
inputs=gr.Textbox(lines=2, placeholder="Enter Caption Here..."), | |
outputs=gr.Image(type="numpy"), | |
title="CUHK Shenzhen Building Text-to-Image Generation", | |
description="Enter a caption of some specific building in CUHK-Shenzhen to generate an image..") | |
iface.launch() | |