DatasetEx.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using MyCaffe.basecode;
using System.Diagnostics;
using MyCaffe.basecode.descriptors;
 
namespace MyCaffe.db.image
{
    public class DatasetEx : IDisposable
    {
        CryptoRandom m_random = null;
        object m_syncObj = new object();
        DatasetFactory m_factory = null;
        DatasetDescriptor m_ds = null;
        ImageSet m_TestingImages = null;
        ImageSet m_TrainingImages = null;
        bool m_bUseTrainingImagesForTesting = false;
        int m_nLastTestingImageIdx = 0;
        int m_nLastTrainingImageIdx = 0;
        List<Guid> m_rgUsers = new List<Guid>();
        int m_nOriginalDsId = 0;
 
        public event EventHandler<CalculateImageMeanArgs> OnCalculateImageMean;
 
        public DatasetEx(Guid user, DatasetFactory factory, CryptoRandom random)
        {
            m_random = random;
 
            if (user != Guid.Empty)
                m_rgUsers.Add(user);
 
            m_factory = new DatasetFactory(factory);
        }
 
        public void Reset()
        {
            m_TestingImages.Reset();
            m_TrainingImages.Reset();
        }
 
        public int AddUser(Guid user)
        {
            m_rgUsers.Add(user);
            return m_rgUsers.Count;
        }
 
        public int RemoveUser(Guid user)
        {
            m_rgUsers.Remove(user);
            return m_rgUsers.Count;
        }
 
        public bool Initialize(DatasetDescriptor ds, WaitHandle[] rgAbort, int nPadW = 0, int nPadH = 0, Log log = null, DB_LOAD_METHOD loadMethod = DB_LOAD_METHOD.LOAD_ALL, int nImageDbLoadLimit = 0, bool bSkipMeanCheck = false)
        {
            lock (m_syncObj)
            {
                if (loadMethod != DB_LOAD_METHOD.LOAD_ALL && nImageDbLoadLimit > 0)
                    throw new Exception("Currently the load-limit only works with the LOAD_ALLL image loading method.");
 
                SimpleDatum imgMean = null;
 
                if (ds != null)
                    m_ds = ds;
 
                if (m_ds.TrainingSource.Width == -1 || m_ds.TrainingSource.Height == -1)
                {
                    log.WriteLine("WARNING: Cannot create a mean image for data sources that contain variable sized images.  The mean check will be skipped.");
                    bSkipMeanCheck = true;
                }
 
                m_TrainingImages = loadImageset("Training", m_ds.TrainingSource, rgAbort, ref imgMean, out m_nLastTrainingImageIdx, nPadW, nPadH, log, loadMethod, nImageDbLoadLimit, m_nLastTrainingImageIdx, (ds == null) ? true : false, bSkipMeanCheck);
                if (m_nLastTrainingImageIdx >= m_ds.TrainingSource.ImageCount)
                    m_nLastTrainingImageIdx = 0;
 
                if (EventWaitHandle.WaitAny(rgAbort, 0) != EventWaitHandle.WaitTimeout)
                    return false;
 
                m_TestingImages = loadImageset("Testing", m_ds.TestingSource, rgAbort, ref imgMean, out m_nLastTestingImageIdx, nPadW, nPadH, log, loadMethod, nImageDbLoadLimit, m_nLastTestingImageIdx, (ds == null) ? true : false, bSkipMeanCheck);
                if (m_nLastTestingImageIdx >= m_ds.TestingSource.ImageCount)
                    m_nLastTestingImageIdx = 0;
 
                if (EventWaitHandle.WaitAny(rgAbort, 0) != EventWaitHandle.WaitTimeout)
                    return false;
 
                return true;
            }
        }
 
        public DatasetEx Clone(bool bReOrganizeByTime = false)
        {
            DatasetEx ds = new DatasetEx(Guid.Empty, m_factory, m_random);
 
            foreach (Guid g in m_rgUsers)
            {
                ds.m_rgUsers.Add(g);
            }
 
            ds.m_ds = new DatasetDescriptor(m_ds);
            ds.m_TestingImages = m_TestingImages.Clone();
            ds.m_TrainingImages = m_TrainingImages.Clone();
            ds.m_bUseTrainingImagesForTesting = m_bUseTrainingImagesForTesting;
 
            if (bReOrganizeByTime)
            {
                ds.m_nOriginalDsId = ds.DatasetID;
                ds.DatasetID *= -1;
 
                int nTestingCount = ds.m_TestingImages.Count;
                int nTrainingCount = ds.m_TrainingImages.Count;
 
                List<SimpleDatum> rgSd = new List<SimpleDatum>();
 
                rgSd.AddRange(ds.m_TestingImages.Images);
                rgSd.AddRange(ds.m_TrainingImages.Images);
 
                int nCount = rgSd.Where(p => p != null).Count();
                if (nCount == 0)
                    throw new Exception("You must first load 'all' images with LOAD_ALL or LOAD_FROM_SERVICE!");
 
                rgSd = rgSd.OrderBy(p => p.Description).ThenBy(p => p.TimeStamp).ToList();
 
                for (int i = 0; i < nTrainingCount; i++)
                {
                    ds.m_TrainingImages.Images[i] = rgSd[i];
                    ds.m_TrainingImages.Images[i].Index = i;
                }
 
                for (int i = 0; i < nTestingCount; i++)
                {
                    ds.m_TestingImages.Images[i] = rgSd[i + nTrainingCount];
                    ds.m_TestingImages.Images[i].Index = i;
                }
 
                ds.m_TrainingImages.Source.ID *= -1;                
                ds.m_TestingImages.Source.ID *= -1;
                ds.Descriptor.TrainingSource.ID *= -1;
                ds.Descriptor.TestingSource.ID *= -1;
 
                ds.m_TrainingImages.ReloadLabelSets();
                ds.m_TestingImages.ReloadLabelSets();
            }
 
            return ds;
        }
 
        public void Relabel(LabelMappingCollection col)
        {
            m_TestingImages.Relabel(col);
            m_TrainingImages.Relabel(col);
        }
 
        public void ReloadLabelSets()
        {
            m_TrainingImages.ReloadLabelSets();
            m_TestingImages.ReloadLabelSets();
        }
 
        public bool UseTrainingImagesForTesting
        {
            get { return m_bUseTrainingImagesForTesting; }
            set { m_bUseTrainingImagesForTesting = value; }
        }
 
        private ImageSet loadImageset(string strType, SourceDescriptor src, WaitHandle[] rgAbort, ref SimpleDatum imgMean, out int nLastImageIdx, int nPadW = 0, int nPadH = 0, Log log = null, DB_LOAD_METHOD loadMethod = DB_LOAD_METHOD.LOAD_ALL, int nImageDbLoadLimit = 0, int nImageDbLoadLimitStartIdx = 0, bool bLoadNext = false, bool bSkipMeanCheck = false)
        {
            try
            {
                RawImageMean imgMeanRaw = null;
               
                m_factory.Open(src);
                nLastImageIdx = nImageDbLoadLimitStartIdx;
 
                if (loadMethod != DB_LOAD_METHOD.LOAD_ALL && !bSkipMeanCheck)
                {
                    if (imgMean == null)
                    {
                        imgMeanRaw = m_factory.GetRawImageMean();
                        if (imgMeanRaw == null)
                        {
                            if (log != null)
                                log.WriteLine("WARNING: No image mean exists in the database, changing image database load from " + loadMethod.ToString() + " to " + DB_LOAD_METHOD.LOAD_ALL.ToString());
 
                            loadMethod = DB_LOAD_METHOD.LOAD_ALL;
                        }
                    }
                }
 
                int nCount = src.ImageCount;
                if (nCount == 0)
                {
                    if (log != null)
                        log.WriteLine("WARNING: Could not find any images with " + strType + " Source = '" + src.Name + "'.  If this is a training dataset, you will need to enable the 'UseTrainingSrcForTesting' setting.");
 
                    if (loadMethod != DB_LOAD_METHOD.LOAD_ON_DEMAND && loadMethod != DB_LOAD_METHOD.LOAD_ON_DEMAND_NOCACHE)
                    {
                        if (log != null)
                            log.WriteLine("Because there are no images in this set, the image loading method has been changed to LOAD_ON_DEMAND for this dataset.");
 
                        loadMethod = DB_LOAD_METHOD.LOAD_ON_DEMAND;
                    }
                }
 
                ImageSet imgset = new ImageSet(m_factory, src, loadMethod, nImageDbLoadLimit, m_random);
 
                if (log != null && nCount > 0)
                    log.WriteLine("Loading '" + src.Name + "' - " + nCount.ToString("N0") + " images.");
 
                if (OnCalculateImageMean != null)
                    imgset.OnCalculateImageMean += OnCalculateImageMean;
 
                if (loadMethod != DB_LOAD_METHOD.LOAD_ON_DEMAND && loadMethod != DB_LOAD_METHOD.LOAD_ON_DEMAND_NOCACHE)
                {
                    bool bDataIsReal = src.IsRealData;
                    int nBatchSize = 20000;
                    Stopwatch sw = new Stopwatch();
 
                    int nImageSize = src.Height * src.Width;
                    if (nImageSize > 60000)
                        nBatchSize = 5000;
                    else if (nBatchSize > 20000)
                        nBatchSize = 7500;
                    else if (nImageSize > 3000)
                        nBatchSize = 10000;
 
                    if (nImageDbLoadLimit <= 0)
                        nImageDbLoadLimit = nCount;
 
                    List<int> rgIdx = (nCount == 0) ? new List<int>() : getIndexList(nImageDbLoadLimitStartIdx, nImageDbLoadLimit);
                    int nIdx = 0;
 
                    sw.Start();
 
                    while (nIdx < rgIdx.Count)
                    {
                        int nImageIdx = rgIdx[nIdx];
                        int nImageCount = Math.Min(rgIdx.Count - nIdx, nBatchSize);
 
                        List<RawImage> rgImg = m_factory.GetRawImagesAt(nImageIdx, nImageCount);
 
                        for (int j = 0; j < rgImg.Count; j++)
                        {
                            SimpleDatum sd1 = m_factory.LoadDatum(rgImg[j], nPadW, nPadH);
                            imgset.Add(nIdx + j, sd1);
 
                            if (sw.Elapsed.TotalMilliseconds > 1000)
                            {
                                if (log != null)
                                {
                                    double dfPct = (double)(nIdx + j) / (double)nCount;
                                    log.Progress = dfPct;
                                    log.WriteLine("image loading at " + dfPct.ToString("P") + "...");
                                }
 
                                sw.Restart();
 
                                if (EventWaitHandle.WaitAny(rgAbort, 0) != EventWaitHandle.WaitTimeout)
                                    return null;
                            }
                        }
 
                        nIdx += rgImg.Count;
 
                        if (loadMethod == DB_LOAD_METHOD.LOAD_ALL && rgImg.Count == 0 && nIdx < nCount)
                        {
                            log.WriteLine("WARNING: Loaded " + nIdx.ToString("N0") + " images, yet " + (nCount - nIdx).ToString("N0") + " images are unaccounted for.  You may need to reindex the dataset.");
                            break;
                        }
                    }
 
                    if (log != null)
                        log.Progress = 0;
 
                    if (rgIdx.Count > 0)
                        nLastImageIdx = rgIdx[rgIdx.Count - 1] + 1;
                }
                else if (bLoadNext)
                {
                    nLastImageIdx += nImageDbLoadLimit;
                }
 
                if (imgMean == null && !bSkipMeanCheck)
                {
                    if (imgMeanRaw == null)
                        imgMeanRaw = m_factory.GetRawImageMean();
 
                    if (imgMeanRaw != null)
                        imgMean = m_factory.LoadDatum(imgMeanRaw, nPadW, nPadH);
                    else
                    {
                        if (log != null)
                            log.WriteLine("Calculating mean...");
 
                        imgMean = imgset.GetImageMean(log, rgAbort);
                        m_factory.PutRawImageMean(imgMean, true);
                    }
                }
 
                if (imgMean != null)
                    imgset.SetImageMean(imgMean);
 
                imgset.CompleteLoad(nLastImageIdx);
 
                return imgset;
            }
            finally
            {
                m_factory.Close();
            }
        }
 
        private List<int> getIndexList(int nStartIdx, int nCount)
        {
            List<int> rgIdx = new List<int>();
 
            for (int i = 0; i < nCount; i++)
            {
                rgIdx.Add(nStartIdx + i);
            }
 
            return rgIdx;
        }
 
        public bool SaveImageMean(int nSrcId, SimpleDatum sd, bool bUpdate)
        {
            if (m_TestingImages.SourceID != nSrcId &&
                m_TrainingImages.SourceID != nSrcId)
                return false;
 
            return m_factory.SaveImageMean(sd, bUpdate, nSrcId);
        }
 
        public SimpleDatum QueryImageMean(int nSrcId)
        {
            if (m_TestingImages.SourceID != nSrcId &&
                m_TrainingImages.SourceID != nSrcId)
                return null;
 
            return m_factory.QueryImageMean(nSrcId);
        }
 
        public void Unload()
        {
            lock (m_syncObj)
            {
                m_TestingImages.Unload();
                m_TrainingImages.Unload();
            }
        }
 
        public double GetPercentageLoaded(out double dfTraining, out double dfTesting)
        {
            int nTrainingTotal = m_TrainingImages.GetTotalCount();
            int nTrainingLoaded = m_TrainingImages.GetLoadedCount();
            int nTestingTotal = m_TestingImages.GetTotalCount();
            int nTestingLoaded = m_TestingImages.GetLoadedCount();
 
            dfTraining = (double)nTrainingLoaded / (double)nTrainingTotal;
            dfTesting = (double)nTestingLoaded / (double)nTestingTotal;
 
            int nTotalLoaded = nTrainingLoaded + nTestingLoaded;
            int nTotalImages = nTrainingTotal + nTestingTotal;
 
            return (double)nTotalLoaded / (double)nTotalImages;
        }
 
        protected virtual void Dispose(bool bDisposing)
        {
            m_ds = null;
 
            if (m_TestingImages != null)
            {
                m_TestingImages.Dispose();
                m_TestingImages = null;
            }
 
            if (m_TrainingImages != null)
            {
                m_TrainingImages.Dispose();
                m_TrainingImages = null;
            }
 
            if (m_factory != null)
            {
                m_factory.Dispose();
                m_factory = null;
            }
        }
 
        public void Dispose()
        {
            Dispose(true);
        }
 
        public ImageSet Find(int nSourceID)
        {
            if (m_TestingImages.SourceID == nSourceID)
            {
                if (m_bUseTrainingImagesForTesting)
                    return m_TrainingImages;
 
                return m_TestingImages;
            }
 
            if (m_TrainingImages.SourceID == nSourceID)
            {
                return m_TrainingImages;
            }
 
            return null;
        }
 
        public ImageSet Find(string strSource)
        {
            if (m_TestingImages.SourceName == strSource)
            {
                if (m_bUseTrainingImagesForTesting)
                    return m_TrainingImages;
 
                return m_TestingImages;
            }
 
            if (m_TrainingImages.SourceName == strSource)
            {
                return m_TrainingImages;
            }
 
            return null;
        }
 
        public DatasetDescriptor Descriptor
        {
            get { return m_ds; }
        }
 
        public int DatasetID
        {
            get { return m_ds.ID; }
            set { m_ds.ID = value; }
        }
 
        public int OriginalDatasetID
        {
            get { return m_nOriginalDsId; }
        }
 
        public string DatasetName
        {
            get { return m_ds.Name; }
        }
    }
}