# All Rights Reserved.
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
# ***************************************************************
__version__ = '1.1.5.6'
from . import lock
with lock.lock_scope():
from . import compiler
from .compiler import LOG, has_cuda
from .compiler import compile_custom_ops, compile_custom_op
import jittor_core as core
from jittor_core import *
from jittor_core.ops import *
from . import compile_extern
from .compile_extern import mkl_ops, mpi, mpi_ops, in_mpi, rank
if core.get_device_count() == 0:
has_cuda = compile_extern.has_cuda = compiler.has_cuda = False
if has_cuda:
from .compile_extern import cudnn, curand, cublas
import contextlib
import numpy as np
from collections import OrderedDict
from collections.abc import Sequence, Mapping
import types
import pickle
import sys
import traceback
class _call_no_record_scope:
def __enter__(self): pass
import numpy as np
import cv2
from PIL import Image
import time
from util.Visualizer import Visualizer
import jittor as jt
from collections import OrderedDict
from options.train_options import TrainOptions
import util.util as util
from models.IS_Model import IS_Model
from jittor.dataset import Dataset
import jittor_core
import os
jt.flags.use_cuda = 1
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
print(jittor_core.get_device_count())
def make_dataset(dir):
images = []
for root, _, fnames in sorted(os.walk(dir)):
for fname in fnames:
path = os.path.join(root, fname)
images.append(path)
return images
class AlignedDataset(Dataset):
def initialize(self, root):
self.part = {
'bg': (0, 0, 512),
