Server_Manga_Translator.py

"""
Manga Translation Server v1.0: YOLO for bubble detection + MangaOCR + mBART for Translation + Colorization - Usage:
pip install fastapi uvicorn ultralytics manga-ocr transformers torch torchvision networkx opencv-python
python Server_Manga_Translator.py
"""
import os, io, base64, time, requests, logging
import numpy as np, cv2, torch, torch.nn as nn, torch.nn.functional as F
from PIL import Image
from fastapi import FastAPI, Body
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from ultralytics import YOLO
from manga_ocr import MangaOcr
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
from torchvision.transforms import ToTensor
from collections import OrderedDict

logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(message)s', datefmt='%H:%M:%S')
log = logging.getLogger(__name__)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_DIR = "./models"; COLOR_DIR = os.path.join(MODEL_DIR, "colorization")
os.makedirs(COLOR_DIR, exist_ok=True)
log.info(f"🚀 {DEVICE}" + (f" ({torch.cuda.get_device_name(0)}, {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB)" if torch.cuda.is_available() else ""))

YOLO_URL="https://huggingface.co/Kirogii/Yolo-Manga_Textbox-Region_Detect/resolve/main/model.pt"
GEN_URL="https://github.com/zyddnys/manga-image-translator/releases/download/beta-0.3/manga-colorization-v2-generator.zip"
DEN_URL="https://github.com/zyddnys/manga-image-translator/releases/download/beta-0.3/manga-colorization-v2-net_rgb.pth"
YOLO_PATH=os.path.join(MODEL_DIR,"yolo.pt")
GEN_PATH=os.path.join(COLOR_DIR,"generator.zip")
DEN_PATH=os.path.join(COLOR_DIR,"net_rgb.pth")

def dl(url, path):
    if os.path.exists(path): return
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(path, 'wb') as f:
            for c in r.iter_content(8192):
                if c: f.write(c)

def ensure(url, path, name):
    if os.path.exists(path): log.info(f"✓ {name} exists")
    else: dl(url, path); log.info(f"↓ {name} downloaded")

ensure(YOLO_URL, YOLO_PATH, "YOLO")
ensure(GEN_URL,  GEN_PATH,  "Generator")
ensure(DEN_URL,  DEN_PATH,  "Denoiser")

yolo = YOLO(YOLO_PATH);
if torch.cuda.is_available(): yolo.to(DEVICE)
log.info("✓ YOLO loaded")
ocr_model = MangaOcr(); log.info("✓ MangaOCR loaded")
mbart_tok = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
mbart_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt").to(DEVICE)
mbart_model.eval(); mbart_tok.src_lang = "ja_XX"; log.info("✓ mBART loaded")

def detect_boxes(im, conf=0.25, iou=0.8, size=640):
    t=time.time()
    res=yolo.predict(im,verbose=False,conf=conf,iou=iou,imgsz=size,device=DEVICE,half=torch.cuda.is_available())[0]
    boxes=[[int(x) for x in b.xyxy[0].cpu().tolist()] for b in res.boxes if (b.xyxy[0][2]-b.xyxy[0][0])>20 and (b.xyxy[0][3]-b.xyxy[0][1])>10]
    dt=time.time()-t; log.info(f"🧭 detect boxes={len(boxes)} in {dt:.3f}s")
    return boxes,dt

def ocr_text(im,box):
    x1,y1,x2,y2=box
    try: return ocr_model(Image.fromarray(im[y1:y2,x1:x2])).strip()
    except: return ""

def translate_batch(texts):
    if not texts: return []
    inp=mbart_tok(texts,return_tensors="pt",padding=True,truncation=True,max_length=256).to(DEVICE)
    with torch.no_grad():
        gen=mbart_model.generate(**inp,forced_bos_token_id=mbart_tok.lang_code_to_id["en_XX"],max_length=128,num_beams=3,early_stopping=True)
    return mbart_tok.batch_decode(gen,skip_special_tokens=True)

_COLOR={'gen':None,'den':None}
def colorize_image(im, denoise=15):
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    if _COLOR['gen'] is None:
        class Selayer(nn.Module):
            def __init__(self,c): super().__init__(); self.global_avgpool=nn.AdaptiveAvgPool2d(1); self.conv1=nn.Conv2d(c,c//16,1,1); self.conv2=nn.Conv2d(c//16,c,1,1); self.relu=nn.ReLU(True); self.sigmoid=nn.Sigmoid()
            def forward(self,x): y=self.global_avgpool(x); y=self.conv1(y); y=self.relu(y); y=self.conv2(y); y=self.sigmoid(y); return x*y
        class BottleneckX_Origin(nn.Module):
            expansion=4
            def __init__(self,inp,p,card,s=1,down=None):
                super().__init__(); self.conv1=nn.Conv2d(inp,p*2,1,bias=False); self.bn1=nn.BatchNorm2d(p*2); self.conv2=nn.Conv2d(p*2,p*2,3,s,1,groups=card,bias=False); self.bn2=nn.BatchNorm2d(p*2); self.conv3=nn.Conv2d(p*2,p*4,1,bias=False); self.bn3=nn.BatchNorm2d(p*4); self.selayer=Selayer(p*4); self.relu=nn.ReLU(True); self.downsample=down
            def forward(self,x): r=x; y=self.conv1(x); y=self.bn1(y); y=self.relu(y); y=self.conv2(y); y=self.bn2(y); y=self.relu(y); y=self.conv3(y); y=self.bn3(y); y=self.selayer(y); r=self.downsample(x) if self.downsample else r; return self.relu(y+r)
        class SEResNeXt_Origin(nn.Module):
            def __init__(self,blk,layers,input_channels=3,cardinality=32):
                super().__init__(); self.cardinality=cardinality; self.inplanes=64
                self.conv1=nn.Conv2d(input_channels,64,7,2,3,bias=False); self.bn1=nn.BatchNorm2d(64); self.relu=nn.ReLU(True)
                self.layer1=self._mk(blk,64,layers[0]); self.layer2=self._mk(blk,128,layers[1],2); self.layer3=self._mk(blk,256,layers[2],2)
            def _mk(self,blk,p,n,s=1):
                down=None
                if s!=1 or self.inplanes!=p*blk.expansion: down=nn.Sequential(nn.Conv2d(self.inplanes,p*blk.expansion,1,s,bias=False),nn.BatchNorm2d(p*blk.expansion))
                L=[blk(self.inplanes,p,self.cardinality,s,down)]; self.inplanes=p*blk.expansion
                for _ in range(1,n): L.append(blk(self.inplanes,p,self.cardinality))
                return nn.Sequential(*L)
            def forward(self,x): x=self.conv1(x); x=self.bn1(x); x1=self.relu(x); x2=self.layer1(x1); x3=self.layer2(x2); x4=self.layer3(x3); return x1,x2,x3,x4
        class ResNeXtBottleneck(nn.Module):
            def __init__(self,c,s=1,card=32,d=1):
                super().__init__(); D=c//2
                self.conv_reduce=nn.Conv2d(c,D,1,bias=False); self.conv_conv=nn.Conv2d(D,D,2+s,s,d,dilation=d,groups=card,bias=False); self.conv_expand=nn.Conv2d(D,c,1,bias=False); self.shortcut=nn.AvgPool2d(2,2) if s!=1 else nn.Identity(); self.selayer=Selayer(c)
            def forward(self,x): b=self.conv_reduce(x); b=F.leaky_relu(b,0.2,True); b=self.conv_conv(b); b=F.leaky_relu(b,0.2,True); b=self.conv_expand(b); b=self.selayer(b); return self.shortcut(x)+b
        class Generator(nn.Module):
            def __init__(self):
                super().__init__(); self.encoder=SEResNeXt_Origin(BottleneckX_Origin,[3,4,6,3],input_channels=1)
                C=lambda i,o,s:nn.Sequential(nn.Conv2d(i,o,3,s,1),nn.LeakyReLU(0.2),nn.Conv2d(o,o,3,1,1),nn.LeakyReLU(0.2))
                self.to0=C(5,32,1); self.to1=C(32,64,2); self.to2=C(64,92,2); self.to3=C(92,128,2); self.to4=C(128,256,2)
                self.deconv_for_decoder=nn.Sequential(nn.ConvTranspose2d(256,128,3,2,1,1),nn.LeakyReLU(0.2),nn.ConvTranspose2d(128,64,3,2,1,1),nn.LeakyReLU(0.2),nn.ConvTranspose2d(64,32,3,1,1,0),nn.LeakyReLU(0.2),nn.ConvTranspose2d(32,3,3,1,1,0),nn.Tanh())
                t4=nn.Sequential(*[ResNeXtBottleneck(512,card=32,d=1) for _ in range(20)])
                self.tunnel4=nn.Sequential(nn.Conv2d(1024+128,512,3,1,1),nn.LeakyReLU(0.2,True),t4,nn.Conv2d(512,1024,3,1,1),nn.PixelShuffle(2),nn.LeakyReLU(0.2,True))
                depth=2
                t3=nn.Sequential(*([ResNeXtBottleneck(256,card=32,d=1) for _ in range(depth)]+[ResNeXtBottleneck(256,card=32,d=2) for _ in range(depth)]+[ResNeXtBottleneck(256,card=32,d=4) for _ in range(depth)]+[ResNeXtBottleneck(256,card=32,d=2),ResNeXtBottleneck(256,card=32,d=1)]))
                self.tunnel3=nn.Sequential(nn.Conv2d(512+256,256,3,1,1),nn.LeakyReLU(0.2,True),t3,nn.Conv2d(256,512,3,1,1),nn.PixelShuffle(2),nn.LeakyReLU(0.2,True))
                t2=nn.Sequential(*([ResNeXtBottleneck(128,card=32,d=1) for _ in range(depth)]+[ResNeXtBottleneck(128,card=32,d=2) for _ in range(depth)]+[ResNeXtBottleneck(128,card=32,d=4) for _ in range(depth)]+[ResNeXtBottleneck(128,card=32,d=2),ResNeXtBottleneck(128,card=32,d=1)]))
                self.tunnel2=nn.Sequential(nn.Conv2d(128+256+64,128,3,1,1),nn.LeakyReLU(0.2,True),t2,nn.Conv2d(128,256,3,1,1),nn.PixelShuffle(2),nn.LeakyReLU(0.2,True))
                t1=nn.Sequential(*[ResNeXtBottleneck(64,card=16,d=1),ResNeXtBottleneck(64,card=16,d=2),ResNeXtBottleneck(64,card=16,d=4),ResNeXtBottleneck(64,card=16,d=2),ResNeXtBottleneck(64,card=16,d=1)])
                self.tunnel1=nn.Sequential(nn.Conv2d(64+32,64,3,1,1),nn.LeakyReLU(0.2,True),t1,nn.Conv2d(64,128,3,1,1),nn.PixelShuffle(2),nn.LeakyReLU(0.2,True))
                self.exit=nn.Sequential(nn.Conv2d(64+32,32,3,1,1),nn.LeakyReLU(0.2,True),nn.Conv2d(32,3,1,1,0))
            def forward(self,sk):
                x0=self.to0(sk); aux=self.to1(x0); aux=self.to2(aux); aux=self.to3(aux)
                x1,x2,x3,x4=self.encoder(sk[:,0:1])
                out=self.tunnel4(torch.cat([x4,aux],1))
                x=self.tunnel3(torch.cat([out,x3],1)); x=self.tunnel2(torch.cat([x,x2,x1],1))
                return torch.tanh(self.exit(torch.cat([x,x0],1)))
        def cat_nm(x, ns):
            N,C,H,W=x.size(); ds=torch.zeros((N,4*C,H//2,W//2),device=x.device,dtype=x.dtype)
            nm=ns.view(N,1,1,1).repeat(1,C,H//2,W//2); idx=((0,0),(0,1),(1,0),(1,1))
            for i,(a,b) in enumerate(idx): ds[:,i:4*C:4]=x[:,:,a::2,b::2]
            return torch.cat((nm,ds),1)
        class IntermediateDnCNN(nn.Module):
            def __init__(self,inp,mid,n):
                super().__init__(); out=4 if inp==5 else 12
                L=[nn.Conv2d(inp,mid,3,1,1,bias=False),nn.ReLU(True)]
                for _ in range(n-2): L+=[nn.Conv2d(mid,mid,3,1,1,bias=False),nn.BatchNorm2d(mid),nn.ReLU(True)]
                L.append(nn.Conv2d(mid,out,3,1,1,bias=False)); self.itermediate_dncnn=nn.Sequential(*L)
            def forward(self,x): return self.itermediate_dncnn(x)
        class FFDNet(nn.Module):
            def __init__(self,ch):
                super().__init__();
                if ch==1: nf,nl,dc=64,15,5
                elif ch==3: nf,nl,dc=96,12,15
                else: raise Exception('bad ch')
                self.intermediate_dncnn=IntermediateDnCNN(dc,nf,nl)
            def forward(self,x,ns):
                h=self.intermediate_dncnn(cat_nm(x.data,ns.data))
                N,Ci,Hi,Wi=h.size(); out=torch.zeros((N,Ci//4,Hi*2,Wi*2),device=h.device,dtype=h.dtype); idx=((0,0),(0,1),(1,0),(1,1))
                for i,(a,b) in enumerate(idx): out[:,:,a::2,b::2]=h[:,i:Ci:4]
                return out
        def resize_pad(img,size=256):
            if len(img.shape)==2: img=np.expand_dims(img,2)
            if img.shape[2]==1: img=np.repeat(img,3,2)
            if img.shape[2]==4: img=img[:,:,:3]
            if img.shape[0]<img.shape[1]:
                r=img.shape[0]/(size*1.5); w=int(np.ceil(img.shape[1]/r)); img=cv2.resize(img,(w,int(size*1.5)),cv2.INTER_AREA); pad=(0, w+(32-w%32)-w); img=np.pad(img,((0,0),(0,pad[1]),(0,0)),'maximum')
            else:
                r=img.shape[1]/size; h=int(np.ceil(img.shape[0]/r)); img=cv2.resize(img,(size,h),cv2.INTER_AREA); pad=(h+(32-h%32)-h, 0); img=np.pad(img,((0,pad[0]),(0,0),(0,0)),'maximum')
            return img[:,:,:1], pad

        gen=Generator().to(dev); gsd=torch.load(GEN_PATH,map_location=dev)
        if isinstance(gsd,dict) and 'state_dict' in gsd: gsd=gsd['state_dict']
        if any(k.startswith('generator.') for k in gsd): gsd={k[len('generator.'):]:v for k,v in gsd.items() if k.startswith('generator.')}
        gen.load_state_dict(gsd, strict=True); gen.eval()
        den=FFDNet(3); dsd=torch.load(DEN_PATH,map_location=dev)
        dsd=OrderedDict((k[7:] if k.startswith('module.') else k,v) for k,v in dsd.items()); den.load_state_dict(dsd); den.eval().to(dev)
        _COLOR['gen'],_COLOR['den']=gen,den

    gen,den=_COLOR['gen'],_COLOR['den']
    sz=min(min(im.shape[:2])//32*32,576)
    if 0<=denoise<=255:
        s=denoise/255; x=im
        if x.ndim<3 or x.shape[2]==1: x=np.repeat(np.expand_dims(x,2),3,2)
        x=x[...,:3]
        if max(x.shape[:2])>1200: r=max(x.shape[:2])/1200; x=cv2.resize(x,(int(x.shape[1]/r),int(x.shape[0]/r)),cv2.INTER_AREA)
        x=x.transpose(2,0,1); x=x/255.0 if x.max()>1.2 else x; x=np.expand_dims(x,0)
        eh,ew=x.shape[2]%2==1,x.shape[3]%2==1
        if eh: x=np.concatenate((x,x[:,:,-1:,:]),2)
        if ew: x=np.concatenate((x,x[:,:,:,-1:]),3)
        xt=torch.tensor(x,device=dev,dtype=torch.float32); ns=torch.tensor([s],device=dev,dtype=torch.float32)
        with torch.no_grad(): out=torch.clamp(xt-den(xt,ns),0,1)
        if eh: out=out[:,:,:-1,:]
        if ew: out=out[:,:,:,:-1]
        im=(out.permute(0,2,3,1).cpu().numpy()[0]*255).astype(np.uint8)
    def resize_pad(img,size=256):
        if len(img.shape)==2: img=np.expand_dims(img,2)
        if img.shape[2]==1: img=np.repeat(img,3,2)
        if img.shape[2]==4: img=img[:,:,:3]
        if img.shape[0]<img.shape[1]:
            r=img.shape[0]/(size*1.5); w=int(np.ceil(img.shape[1]/r)); img=cv2.resize(img,(w,int(size*1.5)),cv2.INTER_AREA); pad=(0, w+(32-w%32)-w); img=np.pad(img,((0,0),(0,pad[1]),(0,0)),'maximum')
        else:
            r=img.shape[1]/size; h=int(np.ceil(img.shape[0]/r)); img=cv2.resize(img,(size,h),cv2.INTER_AREA); pad=(h+(32-h%32)-h, 0); img=np.pad(img,((0,pad[0]),(0,0),(0,0)),'maximum')
        return img[:,:,:1], pad
    img,pad=resize_pad(im,sz)
    cur=ToTensor()(img).unsqueeze(0).to(dev); hint=torch.zeros(1,4,cur.shape[2],cur.shape[3],device=dev)
    with torch.no_grad(): fake=gen(torch.cat([cur,hint],1))
    res=fake[0].cpu().permute(1,2,0)*0.5+0.5
    if pad[0]: res=res[:-pad[0]]
    if pad[1]: res=res[:,:-pad[1]]
    return np.clip(res.numpy()*255,0,255).astype(np.uint8)

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

@app.post("/detect")
async def detect_api(data: dict = Body(...)):
    try:
        im=np.array(Image.open(io.BytesIO(base64.b64decode(data['image']))).convert("RGB"))
        boxes,dt=detect_boxes(im)
        return {'status':'ok','boxes':boxes,'width':im.shape[1],'height':im.shape[0],'time':dt}
    except Exception as e:
        log.error(f"500 detect: {e}"); return JSONResponse(status_code=500, content={'status':'error','message':str(e)})

@app.post("/ocr")
async def ocr_api(data: dict = Body(...)):
    try:
        im=np.array(Image.open(io.BytesIO(base64.b64decode(data['image']))).convert("RGB"))
        texts=[ocr_text(im,b) for b in data.get('boxes',[])]
        return {'status':'ok','texts':texts}
    except Exception as e:
        log.error(f"500 ocr: {e}"); return JSONResponse(status_code=500, content={'status':'error','message':str(e)})

@app.post("/translate_mbart")
async def translate_api(data: dict = Body(...)):
    try:
        t=time.time(); trans=translate_batch(data.get('texts',[])); log.info(f"🌐 translate n={len(trans)} in {time.time()-t:.2f}s")
        return {'status':'ok','translations':trans}
    except Exception as e:
        log.error(f"500 translate: {e}"); return JSONResponse(status_code=500, content={'status':'error','message':str(e)})

@app.post("/colorize")
async def colorize_api(data: dict = Body(...)):
    t=time.time()
    try:
        im=np.array(Image.open(io.BytesIO(base64.b64decode(data['image']))).convert("RGB"))
        d=int(data.get('denoise_sigma',30)); out=colorize_image(im,denoise=d)
        if out.shape[:2]!=im.shape[:2]: out=cv2.resize(out,(im.shape[1],im.shape[0]),cv2.INTER_CUBIC)
        buf=io.BytesIO(); Image.fromarray(out).save(buf,format="PNG")
        dt=time.time()-t; log.info(f"🎨 colorize {im.shape[0]}x{im.shape[1]} σ={d} -> {dt:.2f}s")
        return {'status':'ok','image':base64.b64encode(buf.getvalue()).decode(),'time':dt}
    except Exception as e:
        log.error(f"500 colorize: {e}"); return JSONResponse(status_code=500, content={'status':'error','message':str(e)})

@app.get("/health")
async def health(): return {'status':'ok','device':str(DEVICE)}

if __name__=="__main__":
    import uvicorn
    log.info("🌐 Server: http://localhost:8000")
    uvicorn.run(app,host="0.0.0.0",port=8000,log_level="warning")